use std::collections::HashMap;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_rustls::TlsAcceptor;
use serde::{Deserialize, Serialize};

use remoteingress_protocol::*;

/// Hub configuration.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct HubConfig {
    pub tunnel_port: u16,
    pub target_host: Option<String>,
    #[serde(skip)]
    pub tls_cert_pem: Option<String>,
    #[serde(skip)]
    pub tls_key_pem: Option<String>,
}

impl Default for HubConfig {
    fn default() -> Self {
        Self {
            tunnel_port: 8443,
            target_host: Some("127.0.0.1".to_string()),
            tls_cert_pem: None,
            tls_key_pem: None,
        }
    }
}

/// An allowed edge identity.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AllowedEdge {
    pub id: String,
    pub secret: String,
}

/// Runtime status of a connected edge.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ConnectedEdgeStatus {
    pub edge_id: String,
    pub connected_at: u64,
    pub active_streams: usize,
}

/// Events emitted by the hub.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(tag = "type")]
pub enum HubEvent {
    #[serde(rename_all = "camelCase")]
    EdgeConnected { edge_id: String },
    #[serde(rename_all = "camelCase")]
    EdgeDisconnected { edge_id: String },
    #[serde(rename_all = "camelCase")]
    StreamOpened { edge_id: String, stream_id: u32 },
    #[serde(rename_all = "camelCase")]
    StreamClosed { edge_id: String, stream_id: u32 },
}

/// Hub status response.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct HubStatus {
    pub running: bool,
    pub tunnel_port: u16,
    pub connected_edges: Vec<ConnectedEdgeStatus>,
}

/// The tunnel hub that accepts edge connections and demuxes streams to SmartProxy.
pub struct TunnelHub {
    config: RwLock<HubConfig>,
    allowed_edges: Arc<RwLock<HashMap<String, String>>>, // id -> secret
    connected_edges: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
    event_tx: mpsc::UnboundedSender<HubEvent>,
    event_rx: Mutex<Option<mpsc::UnboundedReceiver<HubEvent>>>,
    shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
    running: RwLock<bool>,
}

struct ConnectedEdgeInfo {
    connected_at: u64,
    active_streams: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
}

impl TunnelHub {
    pub fn new(config: HubConfig) -> Self {
        let (event_tx, event_rx) = mpsc::unbounded_channel();
        Self {
            config: RwLock::new(config),
            allowed_edges: Arc::new(RwLock::new(HashMap::new())),
            connected_edges: Arc::new(Mutex::new(HashMap::new())),
            event_tx,
            event_rx: Mutex::new(Some(event_rx)),
            shutdown_tx: Mutex::new(None),
            running: RwLock::new(false),
        }
    }

    /// Take the event receiver (can only be called once).
    pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<HubEvent>> {
        self.event_rx.lock().await.take()
    }

    /// Update the list of allowed edges.
    pub async fn update_allowed_edges(&self, edges: Vec<AllowedEdge>) {
        let mut map = self.allowed_edges.write().await;
        map.clear();
        for edge in edges {
            map.insert(edge.id, edge.secret);
        }
    }

    /// Get the current hub status.
    pub async fn get_status(&self) -> HubStatus {
        let running = *self.running.read().await;
        let config = self.config.read().await;
        let edges = self.connected_edges.lock().await;

        let mut connected = Vec::new();
        for (id, info) in edges.iter() {
            let streams = info.active_streams.lock().await;
            connected.push(ConnectedEdgeStatus {
                edge_id: id.clone(),
                connected_at: info.connected_at,
                active_streams: streams.len(),
            });
        }

        HubStatus {
            running,
            tunnel_port: config.tunnel_port,
            connected_edges: connected,
        }
    }

    /// Start the hub — listen for TLS connections from edges.
    pub async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let config = self.config.read().await.clone();
        let tls_config = build_tls_config(&config)?;
        let acceptor = TlsAcceptor::from(Arc::new(tls_config));

        let listener = TcpListener::bind(("0.0.0.0", config.tunnel_port)).await?;
        log::info!("Hub listening on port {}", config.tunnel_port);

        let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
        *self.shutdown_tx.lock().await = Some(shutdown_tx);
        *self.running.write().await = true;

        let allowed = self.allowed_edges.clone();
        let connected = self.connected_edges.clone();
        let event_tx = self.event_tx.clone();
        let target_host = config.target_host.unwrap_or_else(|| "127.0.0.1".to_string());

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    result = listener.accept() => {
                        match result {
                            Ok((stream, addr)) => {
                                log::info!("Edge connection from {}", addr);
                                let acceptor = acceptor.clone();
                                let allowed = allowed.clone();
                                let connected = connected.clone();
                                let event_tx = event_tx.clone();
                                let target = target_host.clone();
                                tokio::spawn(async move {
                                    if let Err(e) = handle_edge_connection(
                                        stream, acceptor, allowed, connected, event_tx, target,
                                    ).await {
                                        log::error!("Edge connection error: {}", e);
                                    }
                                });
                            }
                            Err(e) => {
                                log::error!("Accept error: {}", e);
                            }
                        }
                    }
                    _ = shutdown_rx.recv() => {
                        log::info!("Hub shutting down");
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    /// Stop the hub.
    pub async fn stop(&self) {
        if let Some(tx) = self.shutdown_tx.lock().await.take() {
            let _ = tx.send(()).await;
        }
        *self.running.write().await = false;
        // Clear connected edges
        self.connected_edges.lock().await.clear();
    }
}

/// Handle a single edge connection: authenticate, then enter frame loop.
async fn handle_edge_connection(
    stream: TcpStream,
    acceptor: TlsAcceptor,
    allowed: Arc<RwLock<HashMap<String, String>>>,
    connected: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
    event_tx: mpsc::UnboundedSender<HubEvent>,
    target_host: String,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let tls_stream = acceptor.accept(stream).await?;
    let (read_half, write_half) = tokio::io::split(tls_stream);
    let mut buf_reader = BufReader::new(read_half);

    // Read auth line: "EDGE <edgeId> <secret>\n"
    let mut auth_line = String::new();
    buf_reader.read_line(&mut auth_line).await?;
    let auth_line = auth_line.trim();

    let parts: Vec<&str> = auth_line.splitn(3, ' ').collect();
    if parts.len() != 3 || parts[0] != "EDGE" {
        return Err("invalid auth line".into());
    }

    let edge_id = parts[1].to_string();
    let secret = parts[2];

    // Verify credentials
    {
        let edges = allowed.read().await;
        match edges.get(&edge_id) {
            Some(expected) => {
                if !constant_time_eq(secret.as_bytes(), expected.as_bytes()) {
                    return Err(format!("invalid secret for edge {}", edge_id).into());
                }
            }
            None => {
                return Err(format!("unknown edge {}", edge_id).into());
            }
        }
    }

    log::info!("Edge {} authenticated", edge_id);
    let _ = event_tx.send(HubEvent::EdgeConnected {
        edge_id: edge_id.clone(),
    });

    // Track this edge
    let streams: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>> =
        Arc::new(Mutex::new(HashMap::new()));
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();

    {
        let mut edges = connected.lock().await;
        edges.insert(
            edge_id.clone(),
            ConnectedEdgeInfo {
                connected_at: now,
                active_streams: streams.clone(),
            },
        );
    }

    // Shared writer for sending frames back to edge
    let write_half = Arc::new(Mutex::new(write_half));

    // Frame reading loop
    let mut frame_reader = FrameReader::new(buf_reader);

    loop {
        match frame_reader.next_frame().await {
            Ok(Some(frame)) => {
                match frame.frame_type {
                    FRAME_OPEN => {
                        // Payload is PROXY v1 header line
                        let proxy_header = String::from_utf8_lossy(&frame.payload).to_string();

                        // Parse destination port from PROXY header
                        let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443);

                        let stream_id = frame.stream_id;
                        let edge_id_clone = edge_id.clone();
                        let event_tx_clone = event_tx.clone();
                        let streams_clone = streams.clone();
                        let writer_clone = write_half.clone();
                        let target = target_host.clone();

                        let _ = event_tx.send(HubEvent::StreamOpened {
                            edge_id: edge_id.clone(),
                            stream_id,
                        });

                        // Create channel for data from edge to this stream
                        let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(256);
                        {
                            let mut s = streams.lock().await;
                            s.insert(stream_id, data_tx);
                        }

                        // Spawn task: connect to SmartProxy, send PROXY header, pipe data
                        tokio::spawn(async move {
                            let result = async {
                                let mut upstream =
                                    TcpStream::connect((target.as_str(), dest_port)).await?;
                                upstream.write_all(proxy_header.as_bytes()).await?;

                                let (mut up_read, mut up_write) =
                                    upstream.into_split();

                                // Forward data from edge (via channel) to SmartProxy
                                let writer_for_edge_data = tokio::spawn(async move {
                                    while let Some(data) = data_rx.recv().await {
                                        if up_write.write_all(&data).await.is_err() {
                                            break;
                                        }
                                    }
                                    let _ = up_write.shutdown().await;
                                });

                                // Forward data from SmartProxy back to edge
                                let mut buf = vec![0u8; 32768];
                                loop {
                                    match up_read.read(&mut buf).await {
                                        Ok(0) => break,
                                        Ok(n) => {
                                            let frame =
                                                encode_frame(stream_id, FRAME_DATA_BACK, &buf[..n]);
                                            let mut w = writer_clone.lock().await;
                                            if w.write_all(&frame).await.is_err() {
                                                break;
                                            }
                                        }
                                        Err(_) => break,
                                    }
                                }

                                // Send CLOSE_BACK to edge
                                let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]);
                                let mut w = writer_clone.lock().await;
                                let _ = w.write_all(&close_frame).await;

                                writer_for_edge_data.abort();
                                Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
                            }
                            .await;

                            if let Err(e) = result {
                                log::error!("Stream {} error: {}", stream_id, e);
                                // Send CLOSE_BACK on error
                                let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]);
                                let mut w = writer_clone.lock().await;
                                let _ = w.write_all(&close_frame).await;
                            }

                            // Clean up stream
                            {
                                let mut s = streams_clone.lock().await;
                                s.remove(&stream_id);
                            }
                            let _ = event_tx_clone.send(HubEvent::StreamClosed {
                                edge_id: edge_id_clone,
                                stream_id,
                            });
                        });
                    }
                    FRAME_DATA => {
                        let s = streams.lock().await;
                        if let Some(tx) = s.get(&frame.stream_id) {
                            let _ = tx.send(frame.payload).await;
                        }
                    }
                    FRAME_CLOSE => {
                        let mut s = streams.lock().await;
                        s.remove(&frame.stream_id);
                    }
                    _ => {
                        log::warn!("Unexpected frame type {} from edge", frame.frame_type);
                    }
                }
            }
            Ok(None) => {
                log::info!("Edge {} disconnected (EOF)", edge_id);
                break;
            }
            Err(e) => {
                log::error!("Edge {} frame error: {}", edge_id, e);
                break;
            }
        }
    }

    // Cleanup
    {
        let mut edges = connected.lock().await;
        edges.remove(&edge_id);
    }
    let _ = event_tx.send(HubEvent::EdgeDisconnected {
        edge_id: edge_id.clone(),
    });

    Ok(())
}

/// Parse destination port from PROXY v1 header.
fn parse_dest_port_from_proxy(header: &str) -> Option<u16> {
    let parts: Vec<&str> = header.trim().split_whitespace().collect();
    if parts.len() >= 6 {
        parts[5].parse().ok()
    } else {
        None
    }
}

/// Build TLS server config from PEM strings, or auto-generate self-signed.
fn build_tls_config(
    config: &HubConfig,
) -> Result<rustls::ServerConfig, Box<dyn std::error::Error + Send + Sync>> {
    let (cert_pem, key_pem) = match (&config.tls_cert_pem, &config.tls_key_pem) {
        (Some(cert), Some(key)) => (cert.clone(), key.clone()),
        _ => {
            // Generate self-signed certificate
            let cert = rcgen::generate_simple_self_signed(vec!["remoteingress-hub".to_string()])?;
            let cert_pem = cert.cert.pem();
            let key_pem = cert.key_pair.serialize_pem();
            (cert_pem, key_pem)
        }
    };

    let certs = rustls_pemfile_parse_certs(&cert_pem)?;
    let key = rustls_pemfile_parse_key(&key_pem)?;

    let mut config = rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(certs, key)?;

    config.alpn_protocols = vec![b"remoteingress".to_vec()];

    Ok(config)
}

fn rustls_pemfile_parse_certs(
    pem: &str,
) -> Result<Vec<rustls::pki_types::CertificateDer<'static>>, Box<dyn std::error::Error + Send + Sync>>
{
    let mut reader = std::io::Cursor::new(pem.as_bytes());
    let certs = rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?;
    Ok(certs)
}

fn rustls_pemfile_parse_key(
    pem: &str,
) -> Result<rustls::pki_types::PrivateKeyDer<'static>, Box<dyn std::error::Error + Send + Sync>> {
    let mut reader = std::io::Cursor::new(pem.as_bytes());
    let key = rustls_pemfile::private_key(&mut reader)?
        .ok_or("no private key found in PEM")?;
    Ok(key)
}

/// Constant-time comparison of two byte slices.
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff = 0u8;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}