fbc_starter/

server.rs

use axum::{
    http::StatusCode,
    response::Json,
    routing::{get, Router},
    Router as AxumRouter,
};
use serde_json::json;
use std::sync::Arc;
use tokio::signal;
use tower::ServiceBuilder;
use tower_http::compression::CompressionLayer;
use tracing::{info, warn};

use crate::config::Config;
use crate::http::{create_cors_layer, create_trace_layer, health_check, root};
use crate::state::AppState;

#[cfg(feature = "grpc")]
use http_body_util::BodyExt;

/// Web 服务器
pub struct Server {
    config: Config,
    app_state: Arc<AppState>,
    http_router: Option<AxumRouter>,
    #[cfg(feature = "grpc")]
    grpc_router: Option<tonic::transport::server::Router>,
}

/// 服务器构建器，用于在闭包中配置服务器
pub struct ServerBuilder {
    config: Config,
    app_state: Arc<AppState>,
    http_router: Option<AxumRouter>,
    #[cfg(feature = "grpc")]
    grpc_router: Option<tonic::transport::server::Router>,
    #[cfg(feature = "consumer")]
    kafka_handlers: Vec<Arc<dyn crate::messaging::KafkaMessageHandler>>,
}

impl ServerBuilder {
    /// 创建新的服务器构建器
    #[warn(dead_code)]
    fn new(config: Config, app_state: Arc<AppState>) -> Self {
        Self {
            config,
            app_state,
            http_router: None,
            #[cfg(feature = "grpc")]
            grpc_router: None,
            #[cfg(feature = "consumer")]
            kafka_handlers: Vec::new(),
        }
    }

    /// 获取应用状态的引用
    pub fn app_state(&self) -> &Arc<AppState> {
        &self.app_state
    }

    /// 注册 Kafka 消息处理器
    ///
    /// # 参数
    /// - `handler`: 实现了 KafkaMessageHandler trait 的处理器
    ///
    /// # 示例
    /// ```ignore
    /// Server::run(|builder| {
    ///     builder
    ///         .with_kafka_handler(Arc::new(MyHandler::new()))
    ///         .http_router(routes)
    /// })
    /// ```
    #[cfg(feature = "consumer")]
    pub fn with_kafka_handler(
        mut self,
        handler: Arc<dyn crate::messaging::KafkaMessageHandler>,
    ) -> Self {
        self.kafka_handlers.push(handler);
        self
    }

    /// 批量注册 Kafka 消息处理器
    ///
    /// # 示例
    ///
    /// ```no_run
    /// use fbc_starter::Server;
    /// use std::sync::Arc;
    ///
    /// Server::run(|builder| {
    ///     let handlers = vec![
    ///         Arc::new(Handler1::new()) as Arc<dyn fbc_starter::KafkaMessageHandler>,
    ///         Arc::new(Handler2::new()) as Arc<dyn fbc_starter::KafkaMessageHandler>,
    ///     ];
    ///     builder
    ///         .with_kafka_handlers(handlers)
    ///         .http_router(routes)
    /// })
    /// ```
    #[cfg(feature = "consumer")]
    pub fn with_kafka_handlers(
        mut self,
        handlers: Vec<Arc<dyn crate::messaging::KafkaMessageHandler>>,
    ) -> Self {
        self.kafka_handlers.extend(handlers);
        self
    }

    /// 设置 HTTP 路由
    pub fn http_router(mut self, router: AxumRouter) -> Self {
        self.http_router = Some(router);
        self
    }

    /// 设置 gRPC 路由（仅在启用 grpc 特性时可用）
    #[cfg(feature = "grpc")]
    pub fn grpc_router(mut self, router: tonic::transport::server::Router) -> Self {
        self.grpc_router = Some(router);
        self
    }

    /// 获取配置的引用
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// 获取应用状态的引用
    pub fn state(&self) -> &Arc<AppState> {
        &self.app_state
    }

    /// 构建 Server
    fn build(self) -> Server {
        Server {
            config: self.config,
            app_state: self.app_state,
            http_router: self.http_router,
            #[cfg(feature = "grpc")]
            grpc_router: self.grpc_router,
        }
    }
}

impl Server {
    /// 启动服务器（使用闭包配置方式）
    ///
    /// # 参数
    /// - `configure`: 配置闭包，接收 ServerBuilder 用于配置路由
    ///
    /// # 示例
    /// ```rust,no_run
    /// use fbc_starter::{Config, Server};
    /// use axum::{routing::get, Router};
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// Server::run(|builder| {
    ///     // 获取配置
    ///     let config = builder.config();
    ///     
    ///     // 配置 HTTP 路由
    ///     let http_router = Router::new()
    ///         .route("/api/users", get(|| async { "users" }));
    ///     
    ///     // 配置 gRPC 路由（如果启用）
    ///     #[cfg(feature = "grpc")]
    ///     let grpc_router = tonic::transport::Server::builder()
    ///         .add_service(MyServiceServer::new(MyServiceImpl));
    ///     
    ///     builder
    ///         .http_router(http_router)
    ///         #[cfg(feature = "grpc")]
    ///         .grpc_router(grpc_router)
    /// }).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn run<F>(configure: F) -> Result<(), anyhow::Error>
    where
        F: FnOnce(ServerBuilder) -> ServerBuilder,
    {
        // 加载配置
        let config = Config::from_env().unwrap_or_else(|e| {
            eprintln!("⚠ 警告: 无法从环境变量加载配置: {}, 使用默认配置", e);
            Config::default()
        });

        Self::run_with_config(config, configure).await
    }

    /// 使用指定配置启动服务器
    ///
    /// # 参数
    /// - `config`: 服务器配置
    /// - `configure`: 配置闭包，接收 ServerBuilder 用于配置路由
    pub async fn run_with_config<F>(config: Config, configure: F) -> Result<(), anyhow::Error>
    where
        F: FnOnce(ServerBuilder) -> ServerBuilder,
    {
        // 初始化日志
        crate::init_logging(&config).map_err(|e| anyhow::anyhow!("日志初始化失败: {}", e))?;

        tracing::info!(
            "启动 {} v{}",
            env!("CARGO_PKG_NAME"),
            env!("CARGO_PKG_VERSION")
        );

        // 初始化 Nacos（如果配置了且启用了 nacos 特性）
        #[cfg(feature = "nacos")]
        if let Some(ref nacos_config) = config.nacos {
            crate::nacos::init_nacos(nacos_config)
                .await
                .map_err(|e| anyhow::anyhow!("Nacos 初始化失败: {}", e))?;

            // 注册服务
            crate::nacos::register_service(nacos_config, &config.server)
                .await
                .map_err(|e| anyhow::anyhow!("Nacos 服务注册失败: {}", e))?;

            // 订阅服务
            crate::nacos::subscribe_services(nacos_config)
                .await
                .map_err(|e| anyhow::anyhow!("Nacos 服务订阅失败: {}", e))?;

            // 订阅配置
            crate::nacos::subscribe_configs(nacos_config)
                .await
                .map_err(|e| anyhow::anyhow!("Nacos 配置订阅失败: {}", e))?;
        }

        // 创建初始的 builder（不初始化 app_state）
        let mut builder = ServerBuilder::new(config.clone(), Arc::new(AppState::new()));

        // 使用闭包配置服务器
        builder = configure(builder);

        // 初始化应用状态（使用builder中的kafka_handlers）
        #[cfg(feature = "consumer")]
        let app_state = Self::init_app_state(&builder.config, &builder.kafka_handlers).await?;

        #[cfg(not(feature = "consumer"))]
        let app_state = Self::init_app_state(&builder.config).await?;

        // 更新 app_state
        builder.app_state = Arc::new(app_state);

        let server = builder.build();

        // 启动服务器
        server.start_internal().await
    }

    /// 初始化应用状态
    #[cfg(feature = "consumer")]
    async fn init_app_state(
        config: &Config,
        kafka_handlers: &[Arc<dyn crate::messaging::KafkaMessageHandler>],
    ) -> Result<AppState, anyhow::Error> {
        let app_state = AppState::new();

        // 初始化数据库（如果配置了且启用了 database 特性）
        #[cfg(feature = "database")]
        let app_state = if let Some(ref db_config) = config.database {
            let pool = crate::database::init_database(
                &db_config.url,
                db_config.max_connections,
                db_config.min_connections,
            )
            .await
            .map_err(|e| anyhow::anyhow!("数据库初始化失败: {}", e))?;
            app_state.with_database(pool)
        } else {
            app_state
        };

        // 初始化 Redis（如果配置了且启用了 redis 特性）
        #[cfg(feature = "redis")]
        let app_state = if let Some(ref redis_config) = config.redis {
            let conn = crate::cache::redis::init_redis(
                &redis_config.url,
                redis_config.password.as_deref(),
                redis_config.pool_size,
            )
            .await
            .map_err(|e| anyhow::anyhow!("Redis 初始化失败: {}", e))?;
            app_state.with_redis(conn)
        } else {
            app_state
        };

        // Kafka 消息代理初始化（按需初始化 producer 和 consumer）
        #[cfg(feature = "kafka")]
        let app_state = if let Some(ref kafka_config) = config.kafka {
            use std::sync::Arc;
            let mut state = app_state;

            // 按需初始化 Producer（仅当启用 producer 特性且配置了 producer 时）
            #[cfg(feature = "producer")]
            {
                if let Some(ref producer_config) = kafka_config.producer {
                    use crate::messaging::kafka::KafkaProducer;

                    let producer = Arc::new(
                        KafkaProducer::new(&kafka_config.brokers, producer_config)
                            .map_err(|e| anyhow::anyhow!("Kafka 生产者初始化失败: {}", e))?,
                    )
                        as Arc<dyn crate::messaging::MessageProducer + Send + Sync>;
                    state = state.with_message_producer(producer);
                    tracing::info!("✅ Kafka Producer 初始化成功");
                }
            }

            // 按需初始化 Consumer（仅当启用 consumer 特性且配置了 consumer 时）
            #[cfg(feature = "consumer")]
            {
                if let Some(ref consumer_config) = kafka_config.consumer {
                    use crate::messaging::{kafka::KafkaConsumer, KafkaMessageRouter};

                    let consumer = Arc::new(
                        KafkaConsumer::new(&kafka_config.brokers, consumer_config)
                            .map_err(|e| anyhow::anyhow!("Kafka 消费者初始化失败: {}", e))?,
                    )
                        as Arc<dyn crate::messaging::MessageConsumer + Send + Sync>;

                    // 如果有注册的 handlers，自动订阅并设置路由
                    if !kafka_handlers.is_empty() {
                        let router = Arc::new(KafkaMessageRouter::new(kafka_handlers.to_vec()));
                        let subscribe_topics = router.get_subscribe_topics();

                        if !subscribe_topics.is_empty() {
                            let router_clone = router.clone();
                            let handler = Arc::new(move |message: crate::messaging::Message| {
                                let router = router_clone.clone();
                                tokio::spawn(async move {
                                    router.dispatch(message).await;
                                });
                            });

                            consumer
                                .subscribe_topics(subscribe_topics.clone(), handler)
                                .await
                                .map_err(|e| anyhow::anyhow!("Kafka 订阅失败: {}", e))?;

                            tracing::info!(
                                "✅ Kafka Consumer 初始化成功 (group: {}, 自动订阅 {} 个 topics: {:?})",
                                consumer_config.group_id,
                                subscribe_topics.len(),
                                subscribe_topics
                            );
                        } else {
                            tracing::warn!(
                                "⚠️ Kafka Consumer 已初始化，但没有 handler 注册任何 topics"
                            );
                        }
                    } else {
                        tracing::info!(
                            "✅ Kafka Consumer 初始化成功 (group: {}，未注册 handler，需手动订阅)",
                            consumer_config.group_id
                        );
                    }

                    state = state.with_message_consumer(consumer);
                }
            }

            state
        } else {
            app_state
        };

        Ok(app_state)
    }

    /// 初始化应用状态（无 consumer 特性版本）
    #[cfg(not(feature = "consumer"))]
    #[allow(unused_variables)]
    async fn init_app_state(config: &Config) -> Result<AppState, anyhow::Error> {
        let app_state = AppState::new();

        // 初始化数据库（如果配置了且启用了 database 特性）
        #[cfg(feature = "database")]
        let app_state = if let Some(ref db_config) = config.database {
            let pool = crate::database::init_database(
                &db_config.url,
                db_config.max_connections,
                db_config.min_connections,
            )
            .await
            .map_err(|e| anyhow::anyhow!("数据库初始化失败: {}", e))?;
            app_state.with_database(pool)
        } else {
            app_state
        };

        // 初始化 Redis（如果配置了且启用了 redis 特性）
        #[cfg(feature = "redis")]
        let app_state = if let Some(ref redis_config) = config.redis {
            let conn = crate::cache::redis::init_redis(
                &redis_config.url,
                redis_config.password.as_deref(),
                redis_config.pool_size,
            )
            .await
            .map_err(|e| anyhow::anyhow!("Redis 初始化失败: {}", e))?;
            app_state.with_redis(conn)
        } else {
            app_state
        };

        // Kafka Producer 初始化（仅当启用 producer 特性）
        #[cfg(all(feature = "kafka", feature = "producer"))]
        let app_state = if let Some(ref kafka_config) = config.kafka {
            use std::sync::Arc;
            let mut state = app_state;

            if let Some(ref producer_config) = kafka_config.producer {
                use crate::messaging::kafka::KafkaProducer;

                let producer = Arc::new(
                    KafkaProducer::new(&kafka_config.brokers, producer_config)
                        .map_err(|e| anyhow::anyhow!("Kafka 生产者初始化失败: {}", e))?,
                )
                    as Arc<dyn crate::messaging::MessageProducer + Send + Sync>;
                state = state.with_message_producer(producer);
                tracing::info!("✅ Kafka Producer 初始化成功");
            }

            state
        } else {
            app_state
        };

        Ok(app_state)
    }

    /// 创建基础路由（包含健康检查等系统路由）
    fn create_base_router(&self) -> AxumRouter {
        let state = self.app_state.clone();

        Router::new()
            .route("/", get(root))
            .route("/health", get(health_check))
            .fallback(|| async {
                (
                    StatusCode::NOT_FOUND,
                    Json(json!({
                        "error": "Not found",
                        "status": 404
                    })),
                )
            })
            .layer(
                ServiceBuilder::new()
                    .layer(CompressionLayer::new())
                    .into_inner(),
            )
            .layer(create_cors_layer(&self.config))
            .layer(create_trace_layer())
            .with_state(state)
    }

    /// 创建完整路由（合并基础路由、用户自定义路由和 gRPC 服务）
    async fn create_router(&mut self) -> AxumRouter {
        let base_router = self.create_base_router();

        // 合并用户提供的 HTTP 路由
        let router = if let Some(custom) = self.http_router.take() {
            base_router.merge(custom)
        } else {
            base_router
        };

        // 如果启用了 gRPC 特性，添加 gRPC 服务
        #[cfg(feature = "grpc")]
        let router = self.add_grpc_services(router).await;

        // 如果配置了上下文路径，则嵌套到该路径下
        if let Some(ref context_path) = self.config.server.context_path {
            use std::borrow::Cow;
            // 确保路径以 / 开头
            let path: Cow<'_, str> = if context_path.starts_with('/') {
                Cow::Borrowed(context_path.as_str())
            } else {
                // 如果用户没有以 / 开头，自动添加并记录警告
                tracing::warn!("上下文路径 '{}' 应该以 '/' 开头，已自动修正", context_path);
                Cow::Owned(format!("/{}", context_path))
            };
            Router::new().nest(path.as_ref(), router)
        } else {
            router
        }
    }

    /// 添加 gRPC 服务到路由（仅在启用 grpc 特性时可用）
    ///
    /// 使用 axum 的 fallback_service 来托管 gRPC 服务
    /// HTTP 路由优先匹配，未匹配的请求由 gRPC 服务处理
    #[cfg(feature = "grpc")]
    async fn add_grpc_services(&mut self, router: AxumRouter) -> AxumRouter {
        if let Some(grpc_router) = self.grpc_router.take() {
            tracing::info!("集成 gRPC 服务到 axum");

            // 将 gRPC Router 转换为 tower::Service
            let grpc_service = grpc_router.into_service();

            // 创建一个适配器，将 axum Body 转换为 tonic Body
            let grpc_adapter =
                tower::service_fn(move |req: axum::http::Request<axum::body::Body>| {
                    let mut grpc_service = grpc_service.clone();
                    async move {
                        // 转换请求体：axum Body -> tonic BoxBody
                        let (parts, body) = req.into_parts();
                        let body = body
                            .map_err(|e| tonic::Status::internal(e.to_string()))
                            .boxed_unsync();
                        let req = axum::http::Request::from_parts(parts, body);

                        // 调用 gRPC 服务
                        use tower::Service;
                        let res = match grpc_service.call(req).await {
                            Ok(r) => r,
                            Err(_) => {
                                // gRPC 服务错误，返回 500
                                return Ok(axum::http::Response::builder()
                                    .status(StatusCode::INTERNAL_SERVER_ERROR)
                                    .body(axum::body::Body::empty())
                                    .unwrap());
                            }
                        };

                        // 转换响应体：tonic BoxBody -> axum Body
                        let (parts, body) = res.into_parts();
                        let body = axum::body::Body::new(body);
                        Ok::<_, std::convert::Infallible>(axum::http::Response::from_parts(
                            parts, body,
                        ))
                    }
                });

            // 使用 fallback_service 托管 gRPC 服务
            // HTTP 路由优先处理，未匹配的请求由 gRPC 服务处理
            router.fallback_service(grpc_adapter)
        } else {
            router
        }
    }

    /// 内部启动方法（被 start_internal 调用）
    async fn start_internal(mut self) -> Result<(), anyhow::Error> {
        let addr = self.config.server.socket_addr()?;
        let app = self.create_router().await;

        let base_url = if let Some(ref context_path) = self.config.server.context_path {
            format!("http://{}{}", addr, context_path)
        } else {
            format!("http://{}", addr)
        };

        info!("服务器启动在 {}", base_url);
        info!("健康检查: {}/health", base_url);

        #[cfg(feature = "grpc")]
        if self.grpc_router.is_some() {
            info!("gRPC 服务已启用");
        }

        let listener = tokio::net::TcpListener::bind(&addr).await?;

        axum::serve(listener, app)
            .with_graceful_shutdown(shutdown_signal())
            .await?;

        Ok(())
    }
}

/// 优雅关闭信号处理
async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c().await.expect("无法安装 Ctrl+C 信号处理器");
        info!("收到 Ctrl+C 信号，开始优雅关闭...");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("无法安装 SIGTERM 信号处理器")
            .recv()
            .await;
        warn!("收到 SIGTERM 信号，开始优雅关闭...");
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    info!("服务器正在关闭...");
}