use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;

use rand::RngExt;
use tokio::io::{AsyncRead, AsyncWrite, BufReader, BufWriter};
use tokio::net::TcpStream;
use tokio::sync::Mutex;

use crate::errors::{ErrorCode, TetrixError};
use crate::protocol::{read_frame, write_frame};
use crate::types::{AuthRequest, AuthResponse, RequestEnvelope, ResponseEnvelope};

const PROTOCOL_VERSION: &str = "1.0";
const DEFAULT_MAX_CONNECTIONS: u32 = 4;
const DEFAULT_TIMEOUT_MS: u64 = 30_000;
const DEFAULT_MAX_FRAME_SIZE: u32 = 16 << 20; // 16 MiB
const DEFAULT_MAX_RETRIES: u32 = 3;
const DEFAULT_BASE_DELAY_MS: u64 = 500;
const DEFAULT_MAX_DELAY_MS: u64 = 30_000;

/// Configuration for connecting to a TetrixAIDb daemon.
#[derive(Debug, Clone, Default)]
pub struct Config {
    pub endpoint: String,
    pub username: String,
    pub password: String,
    pub max_connections: Option<u32>,
    pub timeout_ms: Option<u64>,
    pub max_frame_size: Option<u32>,
    pub tls: Option<TlsConfig>,
    pub retry: Option<RetryConfig>,
    pub health_check_interval_ms: Option<u64>,
}

/// TLS configuration.
#[derive(Debug, Clone)]
pub struct TlsConfig {
    pub cert_file: Option<String>,
    pub insecure_skip_verify: bool,
}

/// Retry configuration with exponential backoff.
#[derive(Debug, Clone)]
pub struct RetryConfig {
    pub max_retries: Option<u32>,
    pub base_delay_ms: Option<u64>,
    pub max_delay_ms: Option<u64>,
}

enum Scheme {
    Tcp,
    Tls,
    Unix,
}

fn parse_endpoint(endpoint: &str) -> Result<(Scheme, String), TetrixError> {
    if let Some(rest) = endpoint.strip_prefix("unix://") {
        Ok((Scheme::Unix, rest.to_string()))
    } else if let Some(rest) = endpoint.strip_prefix("tls://") {
        Ok((Scheme::Tls, rest.to_string()))
    } else if let Some(rest) = endpoint.strip_prefix("tcp://") {
        Ok((Scheme::Tcp, rest.to_string()))
    } else {
        Ok((Scheme::Tcp, endpoint.to_string()))
    }
}

type BoxedStream = Box<dyn AsyncReadWrite + Send + Unpin>;

trait AsyncReadWrite: AsyncRead + AsyncWrite {}
impl<T: AsyncRead + AsyncWrite> AsyncReadWrite for T {}

struct Connection {
    reader: BufReader<tokio::io::ReadHalf<BoxedStream>>,
    writer: BufWriter<tokio::io::WriteHalf<BoxedStream>>,
}

struct PoolInner {
    idle: Vec<Connection>,
}

/// Async client for the TetrixAIDb daemon with built-in connection pooling,
/// automatic auth-on-connect, retry with exponential backoff, and TLS support.
pub struct TetrixClient {
    config: Config,
    pool: Mutex<PoolInner>,
    active: AtomicU32,
    closed: AtomicBool,
    max_connections: u32,
    timeout_ms: u64,
    max_frame_size: u32,
    tls_connector: Option<Arc<tokio_rustls::TlsConnector>>,
    _health_handle: Option<tokio::task::JoinHandle<()>>,
}

impl TetrixClient {
    /// Create a new client. Does not establish connections until the first request.
    pub fn new(config: Config) -> Result<Arc<Self>, TetrixError> {
        let max_connections = config.max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS);
        let timeout_ms = config.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS);
        let max_frame_size = config.max_frame_size.unwrap_or(DEFAULT_MAX_FRAME_SIZE);

        let (scheme, _) = parse_endpoint(&config.endpoint)?;
        let needs_tls = matches!(scheme, Scheme::Tls) || config.tls.is_some();

        let tls_connector = if needs_tls {
            Some(Arc::new(build_tls_connector(&config.tls)?))
        } else {
            None
        };

        let client = Arc::new(Self {
            config: config.clone(),
            pool: Mutex::new(PoolInner { idle: Vec::new() }),
            active: AtomicU32::new(0),
            closed: AtomicBool::new(false),
            max_connections,
            timeout_ms,
            max_frame_size,
            tls_connector,
            _health_handle: None,
        });

        if let Some(interval_ms) = config.health_check_interval_ms {
            if interval_ms > 0 {
                let weak = Arc::downgrade(&client);
                let handle = tokio::spawn(async move {
                    health_check_loop(weak, interval_ms).await;
                });
                // Safety: we need interior mutability for the handle, but since we only
                // set it once at construction, we use a small trick via Arc::get_mut.
                // In practice, the handle is a background task we don't need to join.
                drop(handle);
            }
        }

        Ok(client)
    }

    /// Close the client and all pooled connections.
    pub async fn close(&self) {
        self.closed.store(true, Ordering::SeqCst);
        let mut pool = self.pool.lock().await;
        pool.idle.clear();
    }

    /// Send a request and decode the response payload.
    pub(crate) async fn request<T: serde::de::DeserializeOwned>(
        &self,
        op: &str,
        payload: &impl serde::Serialize,
    ) -> Result<T, TetrixError> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(TetrixError::new(
                ErrorCode::ConnectionError,
                "client is closed",
            ));
        }

        if self.config.retry.is_some() {
            self.request_with_retry(op, payload).await
        } else {
            self.do_request(op, payload).await
        }
    }

    /// Send a request that returns no meaningful payload (discard response body).
    #[allow(dead_code)]
    pub(crate) async fn request_void(
        &self,
        op: &str,
        payload: &impl serde::Serialize,
    ) -> Result<(), TetrixError> {
        let _: serde_json::Value = self.request(op, payload).await?;
        Ok(())
    }

    async fn do_request<T: serde::de::DeserializeOwned>(
        &self,
        op: &str,
        payload: &impl serde::Serialize,
    ) -> Result<T, TetrixError> {
        let mut conn = self.get_conn().await?;
        let result = self.execute_request(&mut conn, op, payload).await;

        match &result {
            Ok(_) => self.put_conn(conn, true).await,
            Err(e) => {
                let healthy = !matches!(
                    e.code,
                    ErrorCode::ConnectionError | ErrorCode::Timeout
                );
                self.put_conn(conn, healthy).await;
            }
        }

        result
    }

    async fn execute_request<T: serde::de::DeserializeOwned>(
        &self,
        conn: &mut Connection,
        op: &str,
        payload: &impl serde::Serialize,
    ) -> Result<T, TetrixError> {
        let request_id = generate_request_id();
        let envelope = RequestEnvelope {
            protocol_version: PROTOCOL_VERSION.to_string(),
            request_id: request_id.clone(),
            op: op.to_string(),
            idempotency_key: None,
            payload: serde_json::to_value(payload)?,
        };

        let timeout = Duration::from_millis(self.timeout_ms);
        let result = tokio::time::timeout(timeout, async {
            write_frame(&mut conn.writer, &envelope).await?;
            read_frame(&mut conn.reader, Some(self.max_frame_size)).await
        })
        .await;

        let raw_response = match result {
            Ok(Ok(v)) => v,
            Ok(Err(e)) => return Err(e),
            Err(_) => {
                return Err(TetrixError::new(
                    ErrorCode::Timeout,
                    format!("request timed out after {}ms", self.timeout_ms),
                ))
            }
        };

        let response: ResponseEnvelope = serde_json::from_value(raw_response)?;

        if response.status == "error" {
            if let Some(pe) = &response.error {
                return Err(TetrixError::from_protocol(pe));
            }
            return Err(TetrixError::new(
                ErrorCode::InternalError,
                "server returned error with no details",
            ));
        }

        let payload_value = response.payload.unwrap_or(serde_json::Value::Object(
            serde_json::Map::new(),
        ));
        let decoded: T = serde_json::from_value(payload_value).map_err(|e| {
            TetrixError::new(
                ErrorCode::InternalError,
                format!("decode {op} response: {e}"),
            )
        })?;

        Ok(decoded)
    }

    async fn request_with_retry<T: serde::de::DeserializeOwned>(
        &self,
        op: &str,
        payload: &impl serde::Serialize,
    ) -> Result<T, TetrixError> {
        let retry = self.config.retry.as_ref().unwrap();
        let max_retries = retry.max_retries.unwrap_or(DEFAULT_MAX_RETRIES);
        let base_delay = retry.base_delay_ms.unwrap_or(DEFAULT_BASE_DELAY_MS);
        let max_delay = retry.max_delay_ms.unwrap_or(DEFAULT_MAX_DELAY_MS);

        let mut last_err = None;
        for attempt in 0..=max_retries {
            match self.do_request::<T>(op, payload).await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    if !e.is_retryable() || attempt == max_retries {
                        return Err(e);
                    }
                    let delay = compute_backoff_delay(attempt, base_delay, max_delay);
                    tokio::time::sleep(Duration::from_millis(delay)).await;
                    last_err = Some(e);
                }
            }
        }

        Err(last_err.unwrap_or_else(|| {
            TetrixError::new(ErrorCode::InternalError, "retry loop exited unexpectedly")
        }))
    }

    async fn get_conn(&self) -> Result<Connection, TetrixError> {
        {
            let mut pool = self.pool.lock().await;
            if let Some(conn) = pool.idle.pop() {
                self.active.fetch_add(1, Ordering::SeqCst);
                return Ok(conn);
            }
        }

        let current = self.active.load(Ordering::SeqCst);
        if current >= self.max_connections {
            return Err(TetrixError::new(
                ErrorCode::ConnectionError,
                format!(
                    "connection pool exhausted ({current}/{} active)",
                    self.max_connections
                ),
            ));
        }

        self.active.fetch_add(1, Ordering::SeqCst);
        match self.dial().await {
            Ok(conn) => Ok(conn),
            Err(e) => {
                self.active.fetch_sub(1, Ordering::SeqCst);
                Err(e)
            }
        }
    }

    async fn put_conn(&self, conn: Connection, healthy: bool) {
        self.active.fetch_sub(1, Ordering::SeqCst);
        if healthy && !self.closed.load(Ordering::SeqCst) {
            let mut pool = self.pool.lock().await;
            pool.idle.push(conn);
        }
        // unhealthy connections are dropped (closed)
    }

    async fn dial(&self) -> Result<Connection, TetrixError> {
        let (scheme, addr) = parse_endpoint(&self.config.endpoint)?;

        let stream: BoxedStream = match scheme {
            Scheme::Unix => {
                #[cfg(unix)]
                {
                    let stream = tokio::net::UnixStream::connect(&addr).await?;
                    Box::new(stream)
                }
                #[cfg(not(unix))]
                {
                    return Err(TetrixError::new(
                        ErrorCode::InvalidArgument,
                        "unix sockets not supported on this platform",
                    ));
                }
            }
            Scheme::Tcp => {
                let tcp = TcpStream::connect(&addr).await?;
                tcp.set_nodelay(true).ok();
                if let Some(connector) = &self.tls_connector {
                    let server_name = extract_host(&addr)?;
                    let domain = tokio_rustls::rustls::pki_types::ServerName::try_from(
                        server_name.to_string(),
                    )
                    .map_err(|e| {
                        TetrixError::new(
                            ErrorCode::ConnectionError,
                            format!("invalid TLS server name: {e}"),
                        )
                    })?;
                    let tls_stream = connector.connect(domain, tcp).await.map_err(|e| {
                        TetrixError::new(
                            ErrorCode::ConnectionError,
                            format!("TLS handshake failed: {e}"),
                        )
                    })?;
                    Box::new(tls_stream)
                } else {
                    Box::new(tcp)
                }
            }
            Scheme::Tls => {
                let tcp = TcpStream::connect(&addr).await?;
                tcp.set_nodelay(true).ok();
                let connector = self.tls_connector.as_ref().ok_or_else(|| {
                    TetrixError::new(ErrorCode::InternalError, "TLS connector not initialized")
                })?;
                let server_name = extract_host(&addr)?;
                let domain =
                    tokio_rustls::rustls::pki_types::ServerName::try_from(server_name.to_string())
                        .map_err(|e| {
                            TetrixError::new(
                                ErrorCode::ConnectionError,
                                format!("invalid TLS server name: {e}"),
                            )
                        })?;
                let tls_stream = connector.connect(domain, tcp).await.map_err(|e| {
                    TetrixError::new(
                        ErrorCode::ConnectionError,
                        format!("TLS handshake failed: {e}"),
                    )
                })?;
                Box::new(tls_stream)
            }
        };

        let (read_half, write_half) = tokio::io::split(stream);
        let mut conn = Connection {
            reader: BufReader::new(read_half),
            writer: BufWriter::new(write_half),
        };

        self.authenticate(&mut conn).await?;
        Ok(conn)
    }

    async fn authenticate(&self, conn: &mut Connection) -> Result<(), TetrixError> {
        let auth_req = AuthRequest {
            protocol_version: PROTOCOL_VERSION.to_string(),
            op: "AUTH".to_string(),
            username: self.config.username.clone(),
            password: self.config.password.clone(),
        };

        write_frame(&mut conn.writer, &auth_req).await?;
        let raw = read_frame(&mut conn.reader, Some(self.max_frame_size)).await?;
        let auth_resp: AuthResponse = serde_json::from_value(raw)?;

        if auth_resp.status != "ok" {
            if let Some(pe) = &auth_resp.error {
                return Err(TetrixError::from_protocol(pe));
            }
            return Err(TetrixError::new(
                ErrorCode::Unauthenticated,
                format!("auth failed: {}", auth_resp.status),
            ));
        }

        Ok(())
    }
}

fn build_tls_connector(tls_config: &Option<TlsConfig>) -> Result<tokio_rustls::TlsConnector, TetrixError> {
    use tokio_rustls::rustls::{self, ClientConfig};

    let mut root_store = rustls::RootCertStore::empty();

    if let Some(tls) = tls_config {
        if let Some(cert_file) = &tls.cert_file {
            let cert_data = std::fs::read(cert_file).map_err(|e| {
                TetrixError::new(
                    ErrorCode::InvalidArgument,
                    format!("read cert file {cert_file}: {e}"),
                )
            })?;
            let mut cursor = std::io::Cursor::new(cert_data);
            let certs = rustls_pemfile::certs(&mut cursor).collect::<Result<Vec<_>, _>>().map_err(|e| {
                TetrixError::new(
                    ErrorCode::InvalidArgument,
                    format!("parse cert file: {e}"),
                )
            })?;
            for cert in certs {
                root_store.add(cert).map_err(|e| {
                    TetrixError::new(
                        ErrorCode::InvalidArgument,
                        format!("add cert: {e}"),
                    )
                })?;
            }
        }

        if tls.insecure_skip_verify {
            let config = ClientConfig::builder()
                .dangerous()
                .with_custom_certificate_verifier(Arc::new(NoCertVerifier))
                .with_no_client_auth();
            return Ok(tokio_rustls::TlsConnector::from(Arc::new(config)));
        }
    }

    let config = ClientConfig::builder()
        .with_root_certificates(root_store)
        .with_no_client_auth();

    Ok(tokio_rustls::TlsConnector::from(Arc::new(config)))
}

#[derive(Debug)]
struct NoCertVerifier;

impl rustls::client::danger::ServerCertVerifier for NoCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &rustls::pki_types::CertificateDer,
        _intermediates: &[rustls::pki_types::CertificateDer],
        _server_name: &rustls::pki_types::ServerName,
        _ocsp_response: &[u8],
        _now: rustls::pki_types::UnixTime,
    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::danger::ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::pki_types::CertificateDer,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::pki_types::CertificateDer,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        vec![
            rustls::SignatureScheme::RSA_PKCS1_SHA256,
            rustls::SignatureScheme::RSA_PKCS1_SHA384,
            rustls::SignatureScheme::RSA_PKCS1_SHA512,
            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
            rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
            rustls::SignatureScheme::ECDSA_NISTP521_SHA512,
            rustls::SignatureScheme::RSA_PSS_SHA256,
            rustls::SignatureScheme::RSA_PSS_SHA384,
            rustls::SignatureScheme::RSA_PSS_SHA512,
            rustls::SignatureScheme::ED25519,
            rustls::SignatureScheme::ED448,
        ]
    }
}

fn extract_host(addr: &str) -> Result<String, TetrixError> {
    // addr is "host:port" — extract the host part
    if let Some(idx) = addr.rfind(':') {
        Ok(addr[..idx].to_string())
    } else {
        Ok(addr.to_string())
    }
}

fn generate_request_id() -> String {
    let mut rng = rand::rng();
    let bytes: [u8; 8] = rand::RngExt::random(&mut rng);
    hex_encode(&bytes)
}

fn hex_encode(bytes: &[u8]) -> String {
    bytes.iter().map(|b| format!("{b:02x}")).collect()
}

fn compute_backoff_delay(attempt: u32, base_ms: u64, max_ms: u64) -> u64 {
    let mut rng = rand::rng();
    let exp_delay = base_ms.saturating_mul(1u64.wrapping_shl(attempt));
    let capped = exp_delay.min(max_ms);
    let jitter: f64 = rng.random_range(0.0..1.0);
    (capped as f64 * jitter) as u64
}

async fn health_check_loop(client: std::sync::Weak<TetrixClient>, interval_ms: u64) {
    let interval = Duration::from_millis(interval_ms);
    loop {
        tokio::time::sleep(interval).await;
        let Some(client) = client.upgrade() else {
            return;
        };
        if client.closed.load(Ordering::SeqCst) {
            return;
        }
        // Drain idle connections and probe them
        let mut pool = client.pool.lock().await;
        let idle_conns: Vec<Connection> = pool.idle.drain(..).collect();
        drop(pool);

        for conn in idle_conns {
            // For now, we discard idle connections during health check.
            // A production implementation would send a HEALTH frame and keep alive ones.
            drop(conn);
        }
    }
}

use tokio_rustls::rustls;
