use std::convert::Infallible;
use std::net::SocketAddr;

use axum::response::IntoResponse;
use http::Request;
use tonic::body::Body;
use tonic::server::NamedService;
use tonic::transport::Server;
use tonic::transport::server::Router;
use tower::Service;
use tower::layer::util::{Identity, Stack};

use crate::error::InternalError;
use crate::health::HealthService;
use crate::interceptor::Gate;
use crate::reflection::ReflectionService;

type Layer = Stack<oltp::GrpcRequestLog, Identity>;

enum RouterState {
    Pending(Server<Layer>),
    Ready(Router<Layer>),
}

/// 统一装配后的内部 RPC 路由：记录业务 crate 注册的 service 名，
/// health / reflection 在 [`serve`] 时据此生成，避免手工清单漂移。
pub struct InternalRouter {
    state: RouterState,
    services: Vec<&'static str>,
}

impl InternalRouter {
    pub fn add_service<S>(mut self, service: S) -> Self
    where
        S: Service<Request<Body>, Error = Infallible>
            + NamedService
            + Clone
            + Send
            + Sync
            + 'static,
        S::Response: IntoResponse,
        S::Future: Send + 'static,
    {
        self.services.push(S::NAME);
        self.state = match self.state {
            RouterState::Pending(mut server) => RouterState::Ready(server.add_service(service)),
            RouterState::Ready(router) => RouterState::Ready(router.add_service(service)),
        };
        self
    }
}

#[derive(Debug, Clone)]
pub struct ServerConfig {
    pub listen_address: String,
    pub rpc_key: String,
}

impl ServerConfig {
    pub fn new(listen_address: impl Into<String>, rpc_key: impl Into<String>) -> Self {
        Self {
            listen_address: listen_address.into(),
            rpc_key: rpc_key.into(),
        }
    }

    pub fn gate(&self) -> Gate {
        Gate::new(self.rpc_key.clone())
    }

    pub fn listen_address(&self) -> &str {
        &self.listen_address
    }
}

pub fn router(config: &ServerConfig) -> Result<InternalRouter, InternalError> {
    let gate = Gate::new(config.rpc_key.clone());
    if !gate.is_armed() {
        return Err(InternalError::Config(
            "internal rpc key must not be blank".to_string(),
        ));
    }
    Ok(InternalRouter {
        state: RouterState::Pending(Server::builder().layer(oltp::GrpcRequestLog)),
        services: Vec::new(),
    })
}

pub fn parse_address(listen: &str) -> Result<SocketAddr, InternalError> {
    listen
        .parse()
        .map_err(|error| InternalError::Address(format!("{listen}: {error}")))
}

/// 启动内部 RPC 服务端；`shutdown` 完成即优雅停止。
pub async fn serve(
    router: InternalRouter,
    listen: &str,
    shutdown: impl std::future::Future<Output = ()> + Send + 'static,
) -> Result<(), InternalError> {
    let InternalRouter { state, services } = router;
    let health = HealthService::new(&services).server();
    let reflection = ReflectionService::new(&services)?.server();
    let router = match state {
        RouterState::Pending(mut server) => server.add_service(health).add_service(reflection),
        RouterState::Ready(router) => router.add_service(health).add_service(reflection),
    };
    let address = parse_address(listen)?;
    router
        .serve_with_shutdown(address, shutdown)
        .await
        .map_err(InternalError::Transport)
}