use std::collections::BTreeMap;
use std::pin::Pin;
use std::sync::Arc;

use proto::grpc::health::v1::health_check_response::ServingStatus;
use proto::grpc::health::v1::health_server::{Health as HealthTrait, HealthServer};
use proto::grpc::health::v1::{HealthCheckRequest, HealthCheckResponse};
use tonic::{Request, Response, Status};

pub const HEALTH_SERVICE: &str = "grpc.health.v1.Health";
pub const REFLECTION_SERVICE: &str = "grpc.reflection.v1.ServerReflection";

/// 标准 gRPC 健康检查：整体（""）与每个已注册 service 都返回 SERVING。
#[derive(Clone)]
pub struct HealthService {
    statuses: Arc<BTreeMap<String, i32>>,
}

impl HealthService {
    pub fn new(service_names: &[&'static str]) -> Self {
        let mut statuses = BTreeMap::new();
        statuses.insert(String::new(), ServingStatus::Serving as i32);
        for name in service_names {
            statuses.insert((*name).to_string(), ServingStatus::Serving as i32);
        }
        Self {
            statuses: Arc::new(statuses),
        }
    }

    pub fn server(self) -> HealthServer<Self> {
        HealthServer::new(self)
    }

    fn status(&self, service: &str) -> i32 {
        self.statuses
            .get(service)
            .copied()
            .unwrap_or(ServingStatus::ServiceUnknown as i32)
    }
}

#[tonic::async_trait]
impl HealthTrait for HealthService {
    type WatchStream =
        Pin<Box<dyn tokio_stream::Stream<Item = Result<HealthCheckResponse, Status>> + Send>>;

    async fn check(
        &self,
        request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        let service = request.into_inner().service;
        let status = self.status(&service);
        if status == ServingStatus::ServiceUnknown as i32 {
            return Err(Status::not_found("unknown service"));
        }
        Ok(Response::new(HealthCheckResponse { status }))
    }

    async fn watch(
        &self,
        request: Request<HealthCheckRequest>,
    ) -> Result<Response<Self::WatchStream>, Status> {
        let service = request.into_inner().service;
        let status = self.status(&service);
        if status == ServingStatus::ServiceUnknown as i32 {
            return Err(Status::not_found("unknown service"));
        }
        let (tx, rx) = tokio::sync::mpsc::channel(1);
        tokio::spawn(async move {
            let _ = tx.send(Ok(HealthCheckResponse { status })).await;
            // 状态不变：保持流打开直到客户端取消。
            futures_util::future::pending::<()>().await;
        });
        Ok(Response::new(Box::pin(
            tokio_stream::wrappers::ReceiverStream::new(rx),
        )))
    }
}