use crate::codec::DynamicCodec;
use crate::error::Error::GenericError;
use crate::error::Result;
use crate::reflection::{
    fill_pool_from_files, fill_pool_from_reflection, method_desc_to_path,
    reflect_types_for_dynamic_message, reflect_types_for_message,
};
use crate::transport::get_transport;
use crate::{MethodDefinition, ServiceDefinition, json_schema};
use hyper_rustls::HttpsConnector;
use hyper_util::client::legacy::Client;
use hyper_util::client::legacy::connect::HttpConnector;
use log::{info, warn};
pub use prost_reflect::DynamicMessage;
use prost_reflect::ReflectMessage;
use prost_reflect::prost::Message;
use prost_reflect::{DescriptorPool, MethodDescriptor, ServiceDescriptor};
use serde_json::Deserializer;
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::error::Error;
use std::fmt;
use std::fmt::Display;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::ReceiverStream;
use tonic::body::BoxBody;
use tonic::metadata::{MetadataKey, MetadataValue};
use tonic::transport::Uri;
use tonic::{IntoRequest, IntoStreamingRequest, Request, Response, Status, Streaming};
use yaak_tls::ClientCertificateConfig;

#[derive(Clone)]
pub struct GrpcConnection {
    pool: Arc<RwLock<DescriptorPool>>,
    conn: Client<HttpsConnector<HttpConnector>, BoxBody>,
    pub uri: Uri,
    use_reflection: bool,
    max_message_size: usize,
}

#[derive(Default, Debug)]
pub struct GrpcStreamError {
    pub message: String,
    pub status: Option<Status>,
}

impl Error for GrpcStreamError {}

impl Display for GrpcStreamError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.status {
            Some(status) => write!(f, "[{}] {}", status, self.message),
            None => write!(f, "{}", self.message),
        }
    }
}

impl From<String> for GrpcStreamError {
    fn from(value: String) -> Self {
        GrpcStreamError { message: value.to_string(), status: None }
    }
}

impl From<Status> for GrpcStreamError {
    fn from(s: Status) -> Self {
        GrpcStreamError { message: s.message().to_string(), status: Some(s) }
    }
}

impl GrpcConnection {
    pub async fn method(&self, service: &str, method: &str) -> Result<MethodDescriptor> {
        let service = self.service(service).await?;
        let method = service
            .methods()
            .find(|m| m.name() == method)
            .ok_or(GenericError("Failed to find method".to_string()))?;
        Ok(method)
    }

    async fn service(&self, service: &str) -> Result<ServiceDescriptor> {
        let pool = self.pool.read().await;
        let service = pool
            .get_service_by_name(service)
            .ok_or(GenericError("Failed to find service".to_string()))?;
        Ok(service)
    }

    pub async fn unary(
        &self,
        service: &str,
        method: &str,
        message: &str,
        metadata: &BTreeMap<String, String>,
        client_cert: Option<ClientCertificateConfig>,
    ) -> Result<Response<DynamicMessage>> {
        if self.use_reflection {
            reflect_types_for_message(
                self.pool.clone(),
                &self.uri,
                message,
                metadata,
                client_cert,
                self.max_message_size,
            )
            .await?;
        }
        let method = &self.method(&service, &method).await?;
        let input_message = method.input();

        let mut deserializer = Deserializer::from_str(message);
        let req_message = DynamicMessage::deserialize(input_message, &mut deserializer)?;
        deserializer.end()?;

        let mut client = grpc_client(self.conn.clone(), self.uri.clone(), self.max_message_size);

        let mut req = req_message.into_request();
        decorate_req(metadata, &mut req)?;

        let path = method_desc_to_path(method);
        let codec = DynamicCodec::new(method.clone());
        client.ready().await.map_err(|e| GenericError(format!("Failed to connect: {}", e)))?;

        Ok(client.unary(req, path, codec).await?)
    }

    pub async fn serialize_message(
        &self,
        message: &DynamicMessage,
        metadata: &BTreeMap<String, String>,
        client_cert: Option<ClientCertificateConfig>,
    ) -> Result<String> {
        let message = if self.use_reflection {
            reflect_types_for_dynamic_message(
                self.pool.clone(),
                &self.uri,
                message,
                metadata,
                client_cert,
                self.max_message_size,
            )
            .await?;

            let message_name = message.descriptor().full_name().to_string();
            let message_desc = {
                let pool = self.pool.read().await;
                pool.get_message_by_name(&message_name)
                    .ok_or(GenericError(format!("Failed to find message {message_name}")))?
            };
            let mut message_with_updated_pool = DynamicMessage::new(message_desc);
            message_with_updated_pool.merge(message.encode_to_vec().as_slice())?;
            Cow::Owned(message_with_updated_pool)
        } else {
            Cow::Borrowed(message)
        };

        crate::serialize_dynamic_message_json(message.as_ref()).map_err(GenericError)
    }

    pub async fn streaming<F>(
        &self,
        service: &str,
        method: &str,
        stream: ReceiverStream<String>,
        metadata: &BTreeMap<String, String>,
        client_cert: Option<ClientCertificateConfig>,
        on_message: F,
    ) -> Result<Response<Streaming<DynamicMessage>>>
    where
        F: Fn(std::result::Result<String, String>) + Send + Sync + Clone + 'static,
    {
        let method = &self.method(&service, &method).await?;
        let mapped_stream = {
            let input_message = method.input();
            let pool = self.pool.clone();
            let uri = self.uri.clone();
            let md = metadata.clone();
            let use_reflection = self.use_reflection.clone();
            let client_cert = client_cert.clone();
            let max_message_size = self.max_message_size;
            stream
                .then(move |json| {
                    let pool = pool.clone();
                    let uri = uri.clone();
                    let input_message = input_message.clone();
                    let md = md.clone();
                    let use_reflection = use_reflection.clone();
                    let client_cert = client_cert.clone();
                    let on_message = on_message.clone();
                    let json_clone = json.clone();
                    async move {
                        if use_reflection {
                            if let Err(e) = reflect_types_for_message(
                                pool,
                                &uri,
                                &json,
                                &md,
                                client_cert,
                                max_message_size,
                            )
                            .await
                            {
                                warn!("Failed to resolve Any types: {e}");
                            }
                        }
                        let mut de = Deserializer::from_str(&json);
                        match DynamicMessage::deserialize(input_message, &mut de) {
                            Ok(m) => {
                                on_message(Ok(json_clone));
                                Some(m)
                            }
                            Err(e) => {
                                warn!("Failed to deserialize message: {e}");
                                on_message(Err(e.to_string()));
                                None
                            }
                        }
                    }
                })
                .filter_map(|x| x)
        };

        let mut client = grpc_client(self.conn.clone(), self.uri.clone(), self.max_message_size);
        let path = method_desc_to_path(method);
        let codec = DynamicCodec::new(method.clone());

        let mut req = mapped_stream.into_streaming_request();
        decorate_req(metadata, &mut req)?;

        client.ready().await.map_err(|e| GenericError(format!("Failed to connect: {}", e)))?;
        Ok(client.streaming(req, path, codec).await?)
    }

    pub async fn client_streaming<F>(
        &self,
        service: &str,
        method: &str,
        stream: ReceiverStream<String>,
        metadata: &BTreeMap<String, String>,
        client_cert: Option<ClientCertificateConfig>,
        on_message: F,
    ) -> Result<Response<DynamicMessage>>
    where
        F: Fn(std::result::Result<String, String>) + Send + Sync + Clone + 'static,
    {
        let method = &self.method(&service, &method).await?;
        let mapped_stream = {
            let input_message = method.input();
            let pool = self.pool.clone();
            let uri = self.uri.clone();
            let md = metadata.clone();
            let use_reflection = self.use_reflection.clone();
            let client_cert = client_cert.clone();
            let max_message_size = self.max_message_size;
            stream
                .then(move |json| {
                    let pool = pool.clone();
                    let uri = uri.clone();
                    let input_message = input_message.clone();
                    let md = md.clone();
                    let use_reflection = use_reflection.clone();
                    let client_cert = client_cert.clone();
                    let on_message = on_message.clone();
                    let json_clone = json.clone();
                    async move {
                        if use_reflection {
                            if let Err(e) = reflect_types_for_message(
                                pool,
                                &uri,
                                &json,
                                &md,
                                client_cert,
                                max_message_size,
                            )
                            .await
                            {
                                warn!("Failed to resolve Any types: {e}");
                            }
                        }
                        let mut de = Deserializer::from_str(&json);
                        match DynamicMessage::deserialize(input_message, &mut de) {
                            Ok(m) => {
                                on_message(Ok(json_clone));
                                Some(m)
                            }
                            Err(e) => {
                                warn!("Failed to deserialize message: {e}");
                                on_message(Err(e.to_string()));
                                None
                            }
                        }
                    }
                })
                .filter_map(|x| x)
        };

        let mut client = grpc_client(self.conn.clone(), self.uri.clone(), self.max_message_size);
        let path = method_desc_to_path(method);
        let codec = DynamicCodec::new(method.clone());

        let mut req = mapped_stream.into_streaming_request();
        decorate_req(metadata, &mut req)?;

        client.ready().await.map_err(|e| GenericError(format!("Failed to connect: {}", e)))?;
        Ok(client
            .client_streaming(req, path, codec)
            .await
            .map_err(|e| GrpcStreamError { message: e.message().to_string(), status: Some(e) })?)
    }

    pub async fn server_streaming(
        &self,
        service: &str,
        method: &str,
        message: &str,
        metadata: &BTreeMap<String, String>,
    ) -> Result<Response<Streaming<DynamicMessage>>> {
        let method = &self.method(&service, &method).await?;
        let input_message = method.input();

        let mut deserializer = Deserializer::from_str(message);
        let req_message = DynamicMessage::deserialize(input_message, &mut deserializer)?;
        deserializer.end()?;

        let mut client = grpc_client(self.conn.clone(), self.uri.clone(), self.max_message_size);

        let mut req = req_message.into_request();
        decorate_req(metadata, &mut req)?;

        let path = method_desc_to_path(method);
        let codec = DynamicCodec::new(method.clone());
        client.ready().await.map_err(|e| GenericError(format!("Failed to connect: {}", e)))?;
        Ok(client.server_streaming(req, path, codec).await?)
    }
}

fn grpc_client(
    conn: Client<HttpsConnector<HttpConnector>, BoxBody>,
    uri: Uri,
    max_message_size: usize,
) -> tonic::client::Grpc<Client<HttpsConnector<HttpConnector>, BoxBody>> {
    tonic::client::Grpc::with_origin(conn, uri)
        .max_decoding_message_size(max_message_size)
        .max_encoding_message_size(max_message_size)
}

fn message_size_limit(setting: i32) -> usize {
    match setting.try_into() {
        Ok(0) | Err(_) => usize::MAX,
        Ok(limit) => limit,
    }
}

/// Configuration for GrpcHandle to compile proto files
#[derive(Clone)]
pub struct GrpcConfig {
    /// Path to the protoc include directory (vendored/protoc/include)
    pub protoc_include_dir: PathBuf,
    /// Path to the yaakprotoc sidecar binary
    pub protoc_bin_path: PathBuf,
}

pub struct GrpcHandle {
    config: GrpcConfig,
    pools: BTreeMap<String, DescriptorPool>,
}

impl GrpcHandle {
    pub fn new(config: GrpcConfig) -> Self {
        let pools = BTreeMap::new();
        Self { pools, config }
    }
}

impl GrpcHandle {
    /// Remove cached descriptor pool for the given key, if present.
    pub fn invalidate_pool(&mut self, id: &str, uri: &str, proto_files: &Vec<PathBuf>) {
        let key = make_pool_key(id, uri, proto_files);
        self.pools.remove(&key);
    }

    pub async fn reflect(
        &mut self,
        id: &str,
        uri: &str,
        proto_files: &Vec<PathBuf>,
        metadata: &BTreeMap<String, String>,
        validate_certificates: bool,
        client_cert: Option<ClientCertificateConfig>,
        request_message_size: i32,
    ) -> Result<bool> {
        let server_reflection = proto_files.is_empty();
        let key = make_pool_key(id, uri, proto_files);

        // If we already have a pool for this key, reuse it and avoid re-reflection
        if self.pools.contains_key(&key) {
            return Ok(server_reflection);
        }

        let pool = if server_reflection {
            let full_uri = uri_from_str(uri)?;
            fill_pool_from_reflection(
                &full_uri,
                metadata,
                validate_certificates,
                client_cert,
                message_size_limit(request_message_size),
            )
            .await
        } else {
            fill_pool_from_files(&self.config, proto_files).await
        }?;

        self.pools.insert(key, pool.clone());
        Ok(server_reflection)
    }

    pub async fn services(
        &mut self,
        id: &str,
        uri: &str,
        proto_files: &Vec<PathBuf>,
        metadata: &BTreeMap<String, String>,
        validate_certificates: bool,
        client_cert: Option<ClientCertificateConfig>,
        request_message_size: i32,
    ) -> Result<Vec<ServiceDefinition>> {
        // Ensure we have a pool; reflect only if missing
        if self.get_pool(id, uri, proto_files).is_none() {
            info!("Reflecting gRPC services for {} at {}", id, uri);
            self.reflect(
                id,
                uri,
                proto_files,
                metadata,
                validate_certificates,
                client_cert,
                request_message_size,
            )
            .await?;
        }

        let pool = self
            .get_pool(id, uri, proto_files)
            .ok_or(GenericError("Failed to get pool".to_string()))?;
        Ok(self.services_from_pool(&pool))
    }

    fn services_from_pool(&self, pool: &DescriptorPool) -> Vec<ServiceDefinition> {
        pool.services()
            .map(|s| {
                let mut def =
                    ServiceDefinition { name: s.full_name().to_string(), methods: vec![] };
                for method in s.methods() {
                    let input_message = method.input();
                    def.methods.push(MethodDefinition {
                        name: method.name().to_string(),
                        server_streaming: method.is_server_streaming(),
                        client_streaming: method.is_client_streaming(),
                        schema: serde_json::to_string_pretty(&json_schema::message_to_json_schema(
                            &pool,
                            input_message,
                        ))
                        .expect("Failed to serialize JSON schema"),
                    })
                }
                def
            })
            .collect::<Vec<_>>()
    }

    pub async fn connect(
        &mut self,
        id: &str,
        uri: &str,
        proto_files: &Vec<PathBuf>,
        metadata: &BTreeMap<String, String>,
        validate_certificates: bool,
        client_cert: Option<ClientCertificateConfig>,
        request_message_size: i32,
    ) -> Result<GrpcConnection> {
        let use_reflection = proto_files.is_empty();
        let max_message_size = message_size_limit(request_message_size);
        if self.get_pool(id, uri, proto_files).is_none() {
            self.reflect(
                id,
                uri,
                proto_files,
                metadata,
                validate_certificates,
                client_cert.clone(),
                request_message_size,
            )
            .await?;
        }
        let pool = self
            .get_pool(id, uri, proto_files)
            .ok_or(GenericError("Failed to get pool".to_string()))?
            .clone();
        let uri = uri_from_str(uri)?;
        let conn = get_transport(validate_certificates, client_cert.clone())?;
        Ok(GrpcConnection {
            pool: Arc::new(RwLock::new(pool)),
            use_reflection,
            conn,
            uri,
            max_message_size,
        })
    }

    fn get_pool(&self, id: &str, uri: &str, proto_files: &Vec<PathBuf>) -> Option<&DescriptorPool> {
        self.pools.get(make_pool_key(id, uri, proto_files).as_str())
    }
}

pub(crate) fn decorate_req<T>(
    metadata: &BTreeMap<String, String>,
    req: &mut Request<T>,
) -> Result<()> {
    for (k, v) in metadata {
        req.metadata_mut()
            .insert(MetadataKey::from_str(k.as_str())?, MetadataValue::from_str(v.as_str())?);
    }
    Ok(())
}

fn uri_from_str(uri_str: &str) -> Result<Uri> {
    match Uri::from_str(uri_str) {
        Ok(uri) => Ok(uri),
        Err(err) => {
            // Uri::from_str basically only returns "invalid format" so we add more context here
            Err(GenericError(format!("Failed to parse URL, {}", err.to_string())))
        }
    }
}

fn make_pool_key(id: &str, uri: &str, proto_files: &Vec<PathBuf>) -> String {
    let pool_key = format!(
        "{}::{}::{}",
        id,
        uri,
        proto_files
            .iter()
            .map(|p| p.to_string_lossy().to_string())
            .collect::<Vec<String>>()
            .join(":")
    );

    format!("{:x}", md5::compute(pool_key))
}
