#![allow(dead_code)]

use std::sync::Arc;

use anyhow::{Context, Result};
use rustls::{
    pki_types::{PrivateKeyDer, ServerName},
    ClientConfig, RootCertStore,
};
use rustls_pemfile::certs;
use std::fs;
use std::io::BufReader;
use std::pin::Pin;
use std::task::{Context as TaskContext, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;

use crate::discovery::ServiceDiscovery;

pub enum MtlsStream {
    Plain(TcpStream),
    Tls(Box<tokio_rustls::client::TlsStream<TcpStream>>),
}

impl AsyncRead for MtlsStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            MtlsStream::Plain(stream) => Pin::new(stream).poll_read(cx, buf),
            MtlsStream::Tls(stream) => Pin::new(stream.as_mut()).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for MtlsStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
        data: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        match self.get_mut() {
            MtlsStream::Plain(stream) => Pin::new(stream).poll_write(cx, data),
            MtlsStream::Tls(stream) => Pin::new(stream.as_mut()).poll_write(cx, data),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            MtlsStream::Plain(stream) => Pin::new(stream).poll_flush(cx),
            MtlsStream::Tls(stream) => Pin::new(stream.as_mut()).poll_flush(cx),
        }
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            MtlsStream::Plain(stream) => Pin::new(stream).poll_shutdown(cx),
            MtlsStream::Tls(stream) => Pin::new(stream.as_mut()).poll_shutdown(cx),
        }
    }
}

pub struct MtlsClient {
    discovery: Arc<ServiceDiscovery>,
    tls_config: Option<Arc<ClientConfig>>,
}

impl MtlsClient {
    pub fn new(discovery: Arc<ServiceDiscovery>) -> Self {
        Self {
            discovery,
            tls_config: None,
        }
    }

    pub fn with_tls_config(mut self, config: Arc<ClientConfig>) -> Self {
        self.tls_config = Some(config);
        self
    }

    pub async fn connect_to_service(
        &self,
        service_name: &str,
        use_mtls: bool,
    ) -> Result<MtlsStream> {
        let instances = self.discovery.resolve_service(service_name).await?;
        if instances.is_empty() {
            anyhow::bail!("no healthy instances found for service {}", service_name);
        }

        // ラウンドロビン（簡易実装）
        let instance = instances[0].clone();

        let addr = if let Some(mesh_port) = instance.mesh_port {
            format!("{}:{}", instance.ip, mesh_port)
        } else {
            format!("{}:{}", instance.ip, instance.port)
        };

        let stream = TcpStream::connect(&addr).await?;

        if use_mtls {
            let config = self
                .tls_config
                .as_ref()
                .context("TLS config missing for mTLS connection")?;
            let connector = TlsConnector::from(config.clone());
            let server_name = ServerName::try_from(service_name.to_string())
                .context("invalid server name for TLS")?;
            let tls_stream = connector.connect(server_name, stream).await?;
            return Ok(MtlsStream::Tls(Box::new(tls_stream)));
        }

        Ok(MtlsStream::Plain(stream))
    }
}

pub fn build_client_config(
    ca_cert_path: Option<&str>,
    client_cert_path: Option<&str>,
    client_key_path: Option<&str>,
) -> Result<Arc<ClientConfig>> {
    let mut roots = RootCertStore::empty();

    if let Some(ca_path) = ca_cert_path {
        let certs = certs(&mut BufReader::new(fs::File::open(ca_path)?))
            .collect::<std::result::Result<Vec<_>, _>>()?;
        roots.add_parsable_certificates(certs);
    } else {
        // システムのルート証明書を使用
        roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    }

    let config = if let (Some(cert_path), Some(key_path)) = (client_cert_path, client_key_path) {
        let certs = certs(&mut BufReader::new(fs::File::open(cert_path)?))
            .collect::<std::result::Result<Vec<_>, _>>()?;
        let mut key_reader = BufReader::new(fs::File::open(key_path)?);
        let keys = rustls_pemfile::pkcs8_private_keys(&mut key_reader)
            .collect::<std::result::Result<Vec<_>, _>>()?;
        let key = keys
            .into_iter()
            .next()
            .context("no private key found in key file")?;

        ClientConfig::builder()
            .with_root_certificates(roots)
            .with_client_auth_cert(certs, PrivateKeyDer::Pkcs8(key))?
    } else {
        ClientConfig::builder()
            .with_root_certificates(roots)
            .with_no_client_auth()
    };

    Ok(Arc::new(config))
}
