use anyhow::Result;
use bytes::BytesMut;
use serde::Deserialize;
use std::sync::Arc;
use tokio::sync::{mpsc, watch, RwLock};
use tracing::{info, error, warn, debug};

use crate::codec::{Frame, FrameCodec, PacketType};
use crate::crypto;
use crate::keepalive::{self, KeepaliveSignal, LinkHealth};
use crate::telemetry::ConnectionQuality;
use crate::transport;
use crate::transport_trait::{self, TransportSink, TransportStream};
use crate::quic_transport;

/// Client configuration (matches TS IVpnClientConfig).
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ClientConfig {
    pub server_url: String,
    pub server_public_key: String,
    pub dns: Option<Vec<String>>,
    pub mtu: Option<u16>,
    pub keepalive_interval_secs: Option<u64>,
    /// Transport type: "websocket" (default) or "quic".
    pub transport: Option<String>,
    /// For QUIC: SHA-256 hash of server certificate (base64) for cert pinning.
    pub server_cert_hash: Option<String>,
}

/// Client statistics.
#[derive(Debug, Clone, Default)]
pub struct ClientStatistics {
    pub bytes_sent: u64,
    pub bytes_received: u64,
    pub packets_sent: u64,
    pub packets_received: u64,
    pub keepalives_sent: u64,
    pub keepalives_received: u64,
}

/// Client connection state.
#[derive(Debug, Clone, PartialEq)]
pub enum ClientState {
    Disconnected,
    Connecting,
    Handshaking,
    Connected,
    Reconnecting,
    Error(String),
}

impl std::fmt::Display for ClientState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Disconnected => write!(f, "disconnected"),
            Self::Connecting => write!(f, "connecting"),
            Self::Handshaking => write!(f, "handshaking"),
            Self::Connected => write!(f, "connected"),
            Self::Reconnecting => write!(f, "reconnecting"),
            Self::Error(e) => write!(f, "error: {}", e),
        }
    }
}

/// The VPN client.
pub struct VpnClient {
    state: Arc<RwLock<ClientState>>,
    stats: Arc<RwLock<ClientStatistics>>,
    assigned_ip: Arc<RwLock<Option<String>>>,
    shutdown_tx: Option<mpsc::Sender<()>>,
    connected_since: Arc<RwLock<Option<std::time::Instant>>>,
    quality_rx: Option<watch::Receiver<ConnectionQuality>>,
    link_health: Arc<RwLock<LinkHealth>>,
}

impl VpnClient {
    pub fn new() -> Self {
        Self {
            state: Arc::new(RwLock::new(ClientState::Disconnected)),
            stats: Arc::new(RwLock::new(ClientStatistics::default())),
            assigned_ip: Arc::new(RwLock::new(None)),
            shutdown_tx: None,
            connected_since: Arc::new(RwLock::new(None)),
            quality_rx: None,
            link_health: Arc::new(RwLock::new(LinkHealth::Degraded)),
        }
    }

    /// Connect to the VPN server.
    pub async fn connect(&mut self, config: ClientConfig) -> Result<String> {
        if *self.state.read().await != ClientState::Disconnected {
            anyhow::bail!("Client is not disconnected");
        }

        *self.state.write().await = ClientState::Connecting;

        let (shutdown_tx, shutdown_rx) = mpsc::channel::<()>(1);
        self.shutdown_tx = Some(shutdown_tx);

        let state = self.state.clone();
        let stats = self.stats.clone();
        let assigned_ip_ref = self.assigned_ip.clone();
        let connected_since = self.connected_since.clone();
        let link_health = self.link_health.clone();

        // Decode server public key
        let server_pub_key = base64::Engine::decode(
            &base64::engine::general_purpose::STANDARD,
            &config.server_public_key,
        )?;

        // Create transport based on configuration
        let (mut sink, mut stream): (Box<dyn TransportSink>, Box<dyn TransportStream>) = {
            let transport_type = config.transport.as_deref().unwrap_or("auto");
            match transport_type {
                "quic" => {
                    let server_addr = &config.server_url; // For QUIC, serverUrl is host:port
                    let cert_hash = config.server_cert_hash.as_deref();
                    let conn = quic_transport::connect_quic(server_addr, cert_hash).await?;
                    let (quic_sink, quic_stream) = quic_transport::open_quic_streams(conn).await?;
                    info!("Connected via QUIC");
                    (Box::new(quic_sink) as Box<dyn TransportSink>,
                     Box::new(quic_stream) as Box<dyn TransportStream>)
                }
                "websocket" => {
                    let ws = transport::connect_to_server(&config.server_url).await?;
                    let (ws_sink, ws_stream) = transport_trait::split_ws(ws);
                    info!("Connected via WebSocket");
                    (Box::new(ws_sink), Box::new(ws_stream))
                }
                _ => {
                    // "auto" (default): try QUIC first, fall back to WebSocket
                    // Extract host:port from the URL for QUIC attempt
                    let quic_addr = extract_host_port(&config.server_url);
                    let cert_hash = config.server_cert_hash.as_deref();

                    if let Some(ref addr) = quic_addr {
                        match tokio::time::timeout(
                            std::time::Duration::from_secs(3),
                            try_quic_connect(addr, cert_hash),
                        ).await {
                            Ok(Ok((quic_sink, quic_stream))) => {
                                info!("Auto: connected via QUIC to {}", addr);
                                (Box::new(quic_sink) as Box<dyn TransportSink>,
                                 Box::new(quic_stream) as Box<dyn TransportStream>)
                            }
                            Ok(Err(e)) => {
                                debug!("Auto: QUIC failed ({}), falling back to WebSocket", e);
                                let ws = transport::connect_to_server(&config.server_url).await?;
                                let (ws_sink, ws_stream) = transport_trait::split_ws(ws);
                                info!("Auto: connected via WebSocket (QUIC unavailable)");
                                (Box::new(ws_sink), Box::new(ws_stream))
                            }
                            Err(_) => {
                                debug!("Auto: QUIC timed out, falling back to WebSocket");
                                let ws = transport::connect_to_server(&config.server_url).await?;
                                let (ws_sink, ws_stream) = transport_trait::split_ws(ws);
                                info!("Auto: connected via WebSocket (QUIC timed out)");
                                (Box::new(ws_sink), Box::new(ws_stream))
                            }
                        }
                    } else {
                        // Can't extract host:port for QUIC, use WebSocket directly
                        let ws = transport::connect_to_server(&config.server_url).await?;
                        let (ws_sink, ws_stream) = transport_trait::split_ws(ws);
                        info!("Connected via WebSocket");
                        (Box::new(ws_sink), Box::new(ws_stream))
                    }
                }
            }
        };

        // Noise NK handshake (client side = initiator)
        *state.write().await = ClientState::Handshaking;
        let mut initiator = crypto::create_initiator(&server_pub_key)?;
        let mut buf = vec![0u8; 65535];

        // -> e, es
        let len = initiator.write_message(&[], &mut buf)?;
        let init_frame = Frame {
            packet_type: PacketType::HandshakeInit,
            payload: buf[..len].to_vec(),
        };
        let mut frame_bytes = BytesMut::new();
        <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, init_frame, &mut frame_bytes)?;
        sink.send_reliable(frame_bytes.to_vec()).await?;

        // <- e, ee
        let resp_msg = match stream.recv_reliable().await? {
            Some(data) => data,
            None => anyhow::bail!("Connection closed during handshake"),
        };

        let mut frame_buf = BytesMut::from(&resp_msg[..]);
        let frame = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf)?
            .ok_or_else(|| anyhow::anyhow!("Incomplete handshake response frame"))?;

        if frame.packet_type != PacketType::HandshakeResp {
            anyhow::bail!("Expected HandshakeResp, got {:?}", frame.packet_type);
        }

        initiator.read_message(&frame.payload, &mut buf)?;
        let mut noise_transport = initiator.into_transport_mode()?;

        // Receive assigned IP info (encrypted)
        let info_msg = match stream.recv_reliable().await? {
            Some(data) => data,
            None => anyhow::bail!("Connection closed before IP info"),
        };

        let mut frame_buf = BytesMut::from(&info_msg[..]);
        let frame = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf)?
            .ok_or_else(|| anyhow::anyhow!("Incomplete IP info frame"))?;

        let len = noise_transport.read_message(&frame.payload, &mut buf)?;
        let ip_info: serde_json::Value = serde_json::from_slice(&buf[..len])?;
        let assigned_ip = ip_info["assignedIp"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("Missing assignedIp in server response"))?
            .to_string();

        *assigned_ip_ref.write().await = Some(assigned_ip.clone());
        *connected_since.write().await = Some(std::time::Instant::now());
        *state.write().await = ClientState::Connected;

        info!("Connected to VPN, assigned IP: {}", assigned_ip);

        // Create adaptive keepalive monitor (use custom interval if configured)
        let ka_config = config.keepalive_interval_secs.map(|secs| {
            let mut cfg = keepalive::AdaptiveKeepaliveConfig::default();
            cfg.degraded_interval = std::time::Duration::from_secs(secs);
            cfg.healthy_interval = std::time::Duration::from_secs(secs * 2);
            cfg.critical_interval = std::time::Duration::from_secs((secs / 3).max(1));
            cfg
        });
        let (monitor, handle) = keepalive::create_keepalive(ka_config);
        self.quality_rx = Some(handle.quality_rx);

        // Spawn the keepalive monitor
        tokio::spawn(monitor.run());

        // Spawn packet forwarding loop
        let assigned_ip_clone = assigned_ip.clone();
        tokio::spawn(client_loop(
            sink,
            stream,
            noise_transport,
            state,
            stats,
            shutdown_rx,
            handle.signal_rx,
            handle.ack_tx,
            link_health,
        ));

        Ok(assigned_ip_clone)
    }

    /// Disconnect from the VPN server.
    pub async fn disconnect(&mut self) -> Result<()> {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(()).await;
        }
        *self.assigned_ip.write().await = None;
        *self.connected_since.write().await = None;
        *self.state.write().await = ClientState::Disconnected;
        self.quality_rx = None;
        info!("Disconnected from VPN");
        Ok(())
    }

    /// Get current status.
    pub async fn get_status(&self) -> serde_json::Value {
        let state = self.state.read().await;
        let ip = self.assigned_ip.read().await;
        let since = self.connected_since.read().await;

        let mut status = serde_json::json!({
            "state": format!("{}", *state),
        });

        if let Some(ref ip) = *ip {
            status["assignedIp"] = serde_json::json!(ip);
        }
        if let Some(instant) = *since {
            status["uptimeSeconds"] = serde_json::json!(instant.elapsed().as_secs());
        }

        status
    }

    /// Get traffic statistics (includes connection quality).
    pub async fn get_statistics(&self) -> serde_json::Value {
        let stats = self.stats.read().await;
        let since = self.connected_since.read().await;
        let uptime = since.map(|s| s.elapsed().as_secs()).unwrap_or(0);
        let health = self.link_health.read().await;

        let mut result = serde_json::json!({
            "bytesSent": stats.bytes_sent,
            "bytesReceived": stats.bytes_received,
            "packetsSent": stats.packets_sent,
            "packetsReceived": stats.packets_received,
            "keepalivesSent": stats.keepalives_sent,
            "keepalivesReceived": stats.keepalives_received,
            "uptimeSeconds": uptime,
        });

        // Include connection quality if available
        if let Some(ref rx) = self.quality_rx {
            let quality = rx.borrow().clone();
            result["quality"] = serde_json::json!({
                "srttMs": quality.srtt_ms,
                "jitterMs": quality.jitter_ms,
                "minRttMs": quality.min_rtt_ms,
                "maxRttMs": quality.max_rtt_ms,
                "lossRatio": quality.loss_ratio,
                "consecutiveTimeouts": quality.consecutive_timeouts,
                "linkHealth": format!("{}", *health),
                "keepalivesSent": quality.keepalives_sent,
                "keepalivesAcked": quality.keepalives_acked,
            });
        }

        result
    }

    /// Get connection quality snapshot.
    pub fn get_connection_quality(&self) -> Option<ConnectionQuality> {
        self.quality_rx.as_ref().map(|rx| rx.borrow().clone())
    }

    /// Get current link health.
    pub async fn get_link_health(&self) -> LinkHealth {
        *self.link_health.read().await
    }
}

/// The main client packet forwarding loop (runs in a spawned task).
async fn client_loop(
    mut sink: Box<dyn TransportSink>,
    mut stream: Box<dyn TransportStream>,
    mut noise_transport: snow::TransportState,
    state: Arc<RwLock<ClientState>>,
    stats: Arc<RwLock<ClientStatistics>>,
    mut shutdown_rx: mpsc::Receiver<()>,
    mut signal_rx: mpsc::Receiver<KeepaliveSignal>,
    ack_tx: mpsc::Sender<()>,
    link_health: Arc<RwLock<LinkHealth>>,
) {
    let mut buf = vec![0u8; 65535];

    loop {
        tokio::select! {
            msg = stream.recv_reliable() => {
                match msg {
                    Ok(Some(data)) => {
                        let mut frame_buf = BytesMut::from(&data[..]);
                        if let Ok(Some(frame)) = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf) {
                            match frame.packet_type {
                                PacketType::IpPacket => {
                                    match noise_transport.read_message(&frame.payload, &mut buf) {
                                        Ok(len) => {
                                            let mut s = stats.write().await;
                                            s.bytes_received += len as u64;
                                            s.packets_received += 1;
                                        }
                                        Err(e) => {
                                            warn!("Decrypt error: {}", e);
                                            *state.write().await = ClientState::Error(e.to_string());
                                            break;
                                        }
                                    }
                                }
                                PacketType::KeepaliveAck => {
                                    stats.write().await.keepalives_received += 1;
                                    // Signal the keepalive monitor that ACK was received
                                    let _ = ack_tx.send(()).await;
                                }
                                PacketType::Disconnect => {
                                    info!("Server sent disconnect");
                                    *state.write().await = ClientState::Disconnected;
                                    break;
                                }
                                _ => {}
                            }
                        }
                    }
                    Ok(None) => {
                        info!("Connection closed");
                        *state.write().await = ClientState::Disconnected;
                        break;
                    }
                    Err(e) => {
                        error!("Transport error: {}", e);
                        *state.write().await = ClientState::Error(e.to_string());
                        break;
                    }
                }
            }
            signal = signal_rx.recv() => {
                match signal {
                    Some(KeepaliveSignal::SendPing(timestamp_ms)) => {
                        // Embed the timestamp in the keepalive payload (8 bytes, big-endian)
                        let ka_frame = Frame {
                            packet_type: PacketType::Keepalive,
                            payload: timestamp_ms.to_be_bytes().to_vec(),
                        };
                        let mut frame_bytes = BytesMut::new();
                        if <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, ka_frame, &mut frame_bytes).is_ok() {
                            if sink.send_reliable(frame_bytes.to_vec()).await.is_err() {
                                warn!("Failed to send keepalive");
                                *state.write().await = ClientState::Disconnected;
                                break;
                            }
                            stats.write().await.keepalives_sent += 1;
                        }
                    }
                    Some(KeepaliveSignal::PeerDead) => {
                        warn!("Peer declared dead by keepalive monitor");
                        *state.write().await = ClientState::Disconnected;
                        break;
                    }
                    Some(KeepaliveSignal::LinkHealthChanged(health)) => {
                        debug!("Link health changed to: {}", health);
                        *link_health.write().await = health;
                    }
                    None => {
                        // Keepalive monitor channel closed
                        break;
                    }
                }
            }
            _ = shutdown_rx.recv() => {
                // Send disconnect frame
                let dc_frame = Frame {
                    packet_type: PacketType::Disconnect,
                    payload: vec![],
                };
                let mut frame_bytes = BytesMut::new();
                if <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, dc_frame, &mut frame_bytes).is_ok() {
                    let _ = sink.send_reliable(frame_bytes.to_vec()).await;
                }
                let _ = sink.close().await;
                *state.write().await = ClientState::Disconnected;
                break;
            }
        }
    }
}

/// Try to connect via QUIC. Returns transport halves on success.
async fn try_quic_connect(
    addr: &str,
    cert_hash: Option<&str>,
) -> Result<(quic_transport::QuicTransportSink, quic_transport::QuicTransportStream)> {
    let conn = quic_transport::connect_quic(addr, cert_hash).await?;
    let (sink, stream) = quic_transport::open_quic_streams(conn).await?;
    Ok((sink, stream))
}

/// Extract host:port from a WebSocket URL for QUIC auto-fallback.
/// e.g. "ws://127.0.0.1:8080" -> Some("127.0.0.1:8080")
///      "wss://vpn.example.com/tunnel" -> Some("vpn.example.com:443")
///      "127.0.0.1:8080" -> Some("127.0.0.1:8080") (already host:port)
fn extract_host_port(url: &str) -> Option<String> {
    if url.starts_with("ws://") || url.starts_with("wss://") {
        // Parse as URL
        let stripped = if url.starts_with("wss://") {
            &url[6..]
        } else {
            &url[5..]
        };
        // Remove path
        let host_port = stripped.split('/').next()?;
        if host_port.contains(':') {
            Some(host_port.to_string())
        } else {
            // Default port
            let default_port = if url.starts_with("wss://") { 443 } else { 80 };
            Some(format!("{}:{}", host_port, default_port))
        }
    } else if url.contains(':') {
        // Already host:port
        Some(url.to_string())
    } else {
        None
    }
}