use std::collections::{BTreeMap, HashSet};
use std::pin::Pin;
use std::sync::Arc;

use prost::Message;
use proto::grpc::reflection::v1::server_reflection_request::MessageRequest;
use proto::grpc::reflection::v1::server_reflection_response::MessageResponse;
use proto::grpc::reflection::v1::server_reflection_server::{
    ServerReflection as ReflectionTrait, ServerReflectionServer,
};
use proto::grpc::reflection::v1::{
    ErrorResponse, FileDescriptorResponse, ListServiceResponse, ServerReflectionRequest,
    ServerReflectionResponse, ServiceResponse,
};
use tonic::{Request, Response, Status};

/// 基于 `lib/proto` 的 FILE_DESCRIPTOR_SET 实现标准 gRPC 反射（v1）。
///
/// 不引入 tonic-reflection：把 descriptor 解码成 FileDescriptorProto 后自行按
/// list_services / file_by_filename / file_containing_symbol 响应。
#[derive(Clone)]
pub struct ReflectionService {
    files: Arc<Vec<prost_types::FileDescriptorProto>>,
    by_name: Arc<BTreeMap<String, usize>>,
    services: Arc<Vec<String>>,
}

impl ReflectionService {
    pub fn new(service_names: &[&'static str]) -> Result<Self, crate::error::InternalError> {
        let descriptor = prost_types::FileDescriptorSet::decode(proto::FILE_DESCRIPTOR_SET)
            .map_err(|error| {
                crate::error::InternalError::Config(format!(
                    "decoding proto descriptor set: {error}"
                ))
            })?;
        let files = descriptor.file;
        let by_name = files
            .iter()
            .enumerate()
            .filter_map(|(index, file)| file.name.clone().map(|name| (name, index)))
            .collect();
        let mut services = Vec::new();
        for name in service_names {
            services.push((*name).to_string());
        }
        for name in [
            crate::health::HEALTH_SERVICE,
            crate::health::REFLECTION_SERVICE,
        ] {
            if !services.iter().any(|service| service == name) {
                services.push(name.to_string());
            }
        }
        Ok(Self {
            files: Arc::new(files),
            by_name: Arc::new(by_name),
            services: Arc::new(services),
        })
    }

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

    fn response(
        &self,
        request: ServerReflectionRequest,
        message_response: MessageResponse,
    ) -> ServerReflectionResponse {
        ServerReflectionResponse {
            valid_host: request.host.clone(),
            original_request: Some(request),
            message_response: Some(message_response),
        }
    }

    fn error(
        &self,
        request: ServerReflectionRequest,
        code: i32,
        message: &str,
    ) -> ServerReflectionResponse {
        self.response(
            request,
            MessageResponse::ErrorResponse(ErrorResponse {
                error_code: code,
                error_message: message.to_string(),
            }),
        )
    }

    fn file_response(
        &self,
        request: ServerReflectionRequest,
        roots: Vec<usize>,
    ) -> ServerReflectionResponse {
        let mut seen = HashSet::new();
        let mut queue = roots;
        let mut encoded = Vec::new();
        while let Some(index) = queue.pop() {
            if !seen.insert(index) {
                continue;
            }
            let Some(file) = self.files.get(index) else {
                continue;
            };
            encoded.push(file.encode_to_vec());
            for dependency in &file.dependency {
                if let Some(next) = self.by_name.get(dependency) {
                    queue.push(*next);
                }
            }
        }
        self.response(
            request,
            MessageResponse::FileDescriptorResponse(FileDescriptorResponse {
                file_descriptor_proto: encoded,
            }),
        )
    }

    fn handle(&self, request: ServerReflectionRequest) -> ServerReflectionResponse {
        match request.message_request.clone() {
            Some(MessageRequest::ListServices(_)) => {
                let services = self
                    .services
                    .iter()
                    .map(|name| ServiceResponse { name: name.clone() })
                    .collect();
                self.response(
                    request,
                    MessageResponse::ListServicesResponse(ListServiceResponse {
                        service: services,
                    }),
                )
            }
            Some(MessageRequest::FileByFilename(name)) => {
                let Some(index) = self.by_name.get(&name).copied() else {
                    return self.error(request, 5, "file not found");
                };
                self.file_response(request, vec![index])
            }
            Some(MessageRequest::FileContainingSymbol(symbol)) => {
                match self.find_file_by_symbol(&symbol) {
                    Some(index) => self.file_response(request, vec![index]),
                    None => self.error(request, 5, "symbol not found"),
                }
            }
            Some(MessageRequest::AllExtensionNumbersOfType(_))
            | Some(MessageRequest::FileContainingExtension(_)) => {
                self.error(request, 12, "extension reflection is not implemented")
            }
            None => self.error(request, 3, "invalid reflection request"),
        }
    }

    fn find_file_by_symbol(&self, symbol: &str) -> Option<usize> {
        let symbol = symbol.trim();
        for (index, file) in self.files.iter().enumerate() {
            if file.package.as_deref() == Some(symbol) {
                return Some(index);
            }
            let package = file.package.as_deref().unwrap_or_default();
            for service in &file.service {
                let Some(name) = service.name.as_deref() else {
                    continue;
                };
                let full = if package.is_empty() {
                    name.to_string()
                } else {
                    format!("{package}.{name}")
                };
                if full == symbol {
                    return Some(index);
                }
                if let Some(method) = symbol
                    .strip_prefix(&full)
                    .and_then(|rest| rest.strip_prefix('.'))
                {
                    if service
                        .method
                        .iter()
                        .any(|item| item.name.as_deref() == Some(method))
                    {
                        return Some(index);
                    }
                }
            }
            for message in &file.message_type {
                if let Some(name) = message.name.as_deref() {
                    let full = if package.is_empty() {
                        name.to_string()
                    } else {
                        format!("{package}.{name}")
                    };
                    if full == symbol {
                        return Some(index);
                    }
                }
            }
        }
        None
    }
}

#[tonic::async_trait]
impl ReflectionTrait for ReflectionService {
    type ServerReflectionInfoStream = Pin<
        Box<
            dyn tokio_stream::Stream<Item = Result<ServerReflectionResponse, Status>>
                + Send
                + 'static,
        >,
    >;

    async fn server_reflection_info(
        &self,
        request: Request<tonic::Streaming<ServerReflectionRequest>>,
    ) -> Result<Response<Self::ServerReflectionInfoStream>, Status> {
        let mut inbound = request.into_inner();
        let this = self.clone();
        let (tx, rx) = tokio::sync::mpsc::channel(8);
        tokio::spawn(async move {
            loop {
                match tokio_stream::StreamExt::next(&mut inbound).await {
                    None => break,
                    Some(Ok(request)) => {
                        if tx.send(Ok(this.handle(request))).await.is_err() {
                            break;
                        }
                    }
                    Some(Err(_)) => break,
                }
            }
        });
        Ok(Response::new(Box::pin(
            tokio_stream::wrappers::ReceiverStream::new(rx),
        )))
    }
}