use std::io::BufReader;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use anyhow::Context;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, lookup_host};
use tokio::sync::Semaphore;
use tokio::time::timeout;
use tokio_rustls::TlsAcceptor;

const MAX_RELAYS: usize = 1024;
const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_HTTP_HEADERS: usize = 100;

use crate::relay;
use relay::tune_socket;
use crate::stats::Stats;
use crate::GeoLevel;

pub async fn run(
    listen: &str,
    cert: Option<&str>,
    key: Option<&str>,
    no_tls: bool,
    verbose: bool,
    geoip_db: Option<&str>,
    geo_level: GeoLevel,
    stats_file: Option<&str>,
    allow_private: bool,
) -> anyhow::Result<()> {
    let listener = TcpListener::bind(listen)
        .await
        .with_context(|| format!("failed to bind to {listen}"))?;

    let acceptor = if no_tls {
        eprintln!("layover: listening on {listen} (no TLS — development only)");
        None
    } else {
        let acceptor = build_tls_acceptor(cert.unwrap(), key.unwrap())?;
        eprintln!("layover: listening on {listen} (TLS)");
        Some(acceptor)
    };

    let semaphore = Arc::new(Semaphore::new(MAX_RELAYS));
    let conn_id = Arc::new(AtomicU64::new(1));
    let stats = Arc::new(Stats::new(geoip_db, geo_level, stats_file)?);

    // Background task: record history + persist every 60s
    {
        let stats = stats.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(std::time::Duration::from_secs(60));
            interval.tick().await; // skip immediate first tick
            loop {
                interval.tick().await;
                stats.tick();
            }
        });
    }

    loop {
        let (tcp_stream, peer_addr) = listener.accept().await?;
        tcp_stream.set_nodelay(true).ok();
        tune_socket(&tcp_stream);
        let acceptor = acceptor.clone();
        let semaphore = semaphore.clone();
        let stats = stats.clone();
        let id = conn_id.fetch_add(1, Ordering::Relaxed);

        tokio::spawn(async move {
            stats.tcp_total.fetch_add(1, Ordering::Relaxed);
            if verbose {
                eprintln!("layover: [#{id} {peer_addr}] accepted");
            }
            if let Err(e) =
                handle_connection(tcp_stream, peer_addr, acceptor, semaphore, id, verbose, &stats, allow_private)
                    .await
            {
                eprintln!("layover: [#{id} {peer_addr}] {e}");
            }
        });
    }
}

async fn handle_connection(
    tcp_stream: TcpStream,
    peer_addr: SocketAddr,
    acceptor: Option<TlsAcceptor>,
    semaphore: Arc<Semaphore>,
    id: u64,
    verbose: bool,
    stats: &Stats,
    allow_private: bool,
) -> anyhow::Result<()> {
    match acceptor {
        Some(acceptor) => {
            let tls_stream = acceptor
                .accept(tcp_stream)
                .await
                .context("TLS handshake failed")?;
            if verbose {
                eprintln!("layover: [#{id} {peer_addr}] TLS handshake complete");
            }
            handle_client(tls_stream, peer_addr, semaphore, id, verbose, stats, allow_private).await
        }
        None => handle_client(tcp_stream, peer_addr, semaphore, id, verbose, stats, allow_private).await,
    }
}

async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
    mut stream: S,
    peer_addr: SocketAddr,
    semaphore: Arc<Semaphore>,
    id: u64,
    verbose: bool,
    stats: &Stats,
    allow_private: bool,
) -> anyhow::Result<()> {
    let started = Instant::now();

    // Client sends first — destination or HTTP request
    let first_line = timeout(HANDSHAKE_TIMEOUT, relay::read_line(&mut stream))
        .await
        .context("handshake timed out")?
        .context("failed to read first line")?;

    let trimmed = first_line.trim();

    // Detect HTTP requests (browsers visiting the server)
    if let Some(path) = parse_http_path(trimmed) {
        stats.http_requests.fetch_add(1, Ordering::Relaxed);
        eprintln!("layover: [#{id} {peer_addr}] HTTP {path}");
        drain_http_headers(&mut stream).await?;
        return match path {
            "/stats" => serve_stats_page(&mut stream).await,
            "/stats.json" => serve_stats_json(&mut stream, stats).await,
            _ => serve_landing_page(&mut stream).await,
        };
    }

    // Protocol: "V2 destination\n" → "OK\n"
    let destination = match trimmed.strip_prefix("V2 ") {
        Some(dest) => dest.to_string(),
        None => {
            relay::write_line(&mut stream, "ERR: unsupported protocol, upgrade layover").await?;
            anyhow::bail!("unsupported protocol from {peer_addr}: {trimmed}");
        }
    };

    stats.relays_total.fetch_add(1, Ordering::Relaxed);

    eprintln!("layover: [#{id} {peer_addr}] -> {destination}");

    // Resolve destination — skip async DNS for raw IP:port to avoid thread pool dispatch
    let addrs: Vec<SocketAddr> = if let Ok(addr) = destination.parse::<SocketAddr>() {
        vec![addr]
    } else {
        match lookup_host(&*destination).await {
            Ok(addrs) => addrs.collect(),
            Err(e) => {
                relay::write_line(&mut stream, "ERR: could not resolve destination").await?;
                eprintln!("layover: [#{id} {peer_addr}] failed to resolve {destination}: {e}");
                stats.relays_failed.fetch_add(1, Ordering::Relaxed);
                return Ok(());
            }
        }
    };

    let allowed: Vec<SocketAddr> = addrs.into_iter().filter(|a| allow_private || !is_blocked_ip(a.ip())).collect();
    if allowed.is_empty() {
        relay::write_line(&mut stream, "ERR: destination blocked").await?;
        eprintln!("layover: [#{id} {peer_addr}] blocked {destination} (private/reserved IP)");
        stats.relays_failed.fetch_add(1, Ordering::Relaxed);
        return Ok(());
    }

    if verbose {
        let ips: Vec<_> = allowed.iter().map(ToString::to_string).collect();
        eprintln!("layover: [#{id} {peer_addr}] resolved {destination} -> {}", ips.join(", "));
    }

    // Acquire relay permit — only the relay phase is guarded
    let _permit = match semaphore.try_acquire_owned() {
        Ok(permit) => permit,
        Err(_) => {
            relay::write_line(&mut stream, "ERR: server busy").await?;
            eprintln!("layover: [#{id} {peer_addr}] rejected (relay limit)");
            stats.relays_failed.fetch_add(1, Ordering::Relaxed);
            return Ok(());
        }
    };

    match TcpStream::connect(allowed.as_slice()).await {
        Ok(mut dest_stream) => {
            dest_stream.set_nodelay(true).ok();
            tune_socket(&dest_stream);
            let dest_ip = dest_stream.peer_addr().ok().map(|a| a.ip());
            if verbose {
                let dest_local = dest_stream.local_addr().ok();
                eprintln!(
                    "layover: [#{id} {peer_addr}] connected {local} -> {remote}",
                    local = dest_local.map_or("?".into(), |a| a.to_string()),
                    remote = dest_ip.map_or("?".into(), |a| a.to_string()),
                );
            }
            stats.record_source_ip(peer_addr.ip());
            if let Some(ip) = dest_ip {
                stats.record_dest_ip(ip);
            }
            stats.relays_active.fetch_add(1, Ordering::Relaxed);
            relay::write_line(&mut stream, "OK").await?;
            let (up, down) = relay::relay(&mut stream, &mut dest_stream).await;
            stats.relays_active.fetch_sub(1, Ordering::Relaxed);
            stats.relays_completed.fetch_add(1, Ordering::Relaxed);
            stats.bytes_up.fetch_add(up, Ordering::Relaxed);
            stats.bytes_down.fetch_add(down, Ordering::Relaxed);
            let elapsed = started.elapsed();
            eprintln!(
                "layover: [#{id} {peer_addr}] closed ({up}B up, {down}B down, {elapsed:.1?})",
            );
        }
        Err(e) => {
            relay::write_line(&mut stream, "ERR: could not connect").await?;
            eprintln!("layover: [#{id} {peer_addr}] failed to reach {destination}: {e}");
            stats.relays_failed.fetch_add(1, Ordering::Relaxed);
        }
    }

    Ok(())
}

/// Parse HTTP request line, return path if valid. e.g. "GET /stats HTTP/1.1" -> Some("/stats")
fn parse_http_path(line: &str) -> Option<&str> {
    let mut parts = line.split_whitespace();
    let method = parts.next()?;
    if !matches!(method, "GET" | "HEAD" | "POST") {
        return None;
    }
    parts.next()
}

async fn drain_http_headers<S: AsyncRead + Unpin>(stream: &mut S) -> anyhow::Result<()> {
    for _ in 0..MAX_HTTP_HEADERS {
        let line = timeout(HANDSHAKE_TIMEOUT, relay::read_line(stream))
            .await
            .context("timed out reading HTTP headers")?
            .context("failed to read HTTP header")?;
        if line.trim().is_empty() {
            break;
        }
    }
    Ok(())
}

async fn serve_http_response<S: AsyncWrite + Unpin>(
    stream: &mut S,
    content_type: &str,
    body: &str,
) -> anyhow::Result<()> {
    let response = format!(
        "HTTP/1.1 200 OK\r\n\
         Content-Type: {content_type}\r\n\
         Content-Length: {}\r\n\
         Connection: close\r\n\
         X-Content-Type-Options: nosniff\r\n\
         X-Frame-Options: DENY\r\n\
         Referrer-Policy: no-referrer\r\n\
         Strict-Transport-Security: max-age=63072000; includeSubDomains\r\n\
         \r\n\
         {body}",
        body.len(),
    );
    stream.write_all(response.as_bytes()).await?;
    stream.flush().await?;
    Ok(())
}

async fn serve_landing_page<S: AsyncWrite + Unpin>(stream: &mut S) -> anyhow::Result<()> {
    let body = include_str!("landing.html").replace("{{VERSION}}", env!("CARGO_PKG_VERSION"));
    serve_http_response(stream, "text/html; charset=utf-8", &body).await
}

async fn serve_stats_page<S: AsyncWrite + Unpin>(stream: &mut S) -> anyhow::Result<()> {
    serve_http_response(stream, "text/html; charset=utf-8", include_str!("stats.html")).await
}

async fn serve_stats_json<S: AsyncWrite + Unpin>(
    stream: &mut S,
    stats: &Stats,
) -> anyhow::Result<()> {
    let json = serde_json::to_string(&stats.snapshot())?;
    let response = format!(
        "HTTP/1.1 200 OK\r\n\
         Content-Type: application/json\r\n\
         Content-Length: {}\r\n\
         Connection: close\r\n\
         Access-Control-Allow-Origin: *\r\n\
         Cache-Control: no-cache\r\n\
         \r\n\
         {json}",
        json.len(),
    );
    stream.write_all(response.as_bytes()).await?;
    stream.flush().await?;
    Ok(())
}


/// Block connections to loopback, private (RFC 1918), link-local, shared
/// address space (RFC 6598), benchmarking, reserved, and IPv4-mapped IPv6.
fn is_blocked_ip(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => is_blocked_v4(&v4),
        IpAddr::V6(v6) => {
            // Check IPv4-mapped IPv6 (::ffff:x.x.x.x) against v4 rules
            if let Some(mapped) = v6.to_ipv4_mapped() {
                return is_blocked_v4(&mapped);
            }
            v6.is_loopback()
                || v6.is_unspecified()
                || (v6.segments()[0] & 0xfe00) == 0xfc00 // unique local (fc00::/7)
                || (v6.segments()[0] & 0xffc0) == 0xfe80 // link-local (fe80::/10)
        }
    }
}

fn is_blocked_v4(v4: &std::net::Ipv4Addr) -> bool {
    let octets = v4.octets();
    v4.is_loopback()                        // 127.0.0.0/8
        || v4.is_private()                  // 10/8, 172.16/12, 192.168/16
        || v4.is_link_local()              // 169.254.0.0/16 (incl. cloud metadata)
        || v4.is_unspecified()             // 0.0.0.0
        || v4.is_broadcast()               // 255.255.255.255
        || (octets[0] == 100 && (octets[1] & 0xc0) == 64)  // 100.64.0.0/10 (RFC 6598)
        || (octets[0] == 198 && (octets[1] & 0xfe) == 18)  // 198.18.0.0/15 (benchmarking)
        || octets[0] >= 240               // 240.0.0.0/4 (reserved)
}

fn build_tls_acceptor(cert_path: &str, key_path: &str) -> anyhow::Result<TlsAcceptor> {
    use tokio_rustls::rustls::{CipherSuite, ServerConfig, version::TLS13};
    use tokio_rustls::rustls::crypto::aws_lc_rs;

    let cert_file = std::fs::File::open(cert_path)
        .with_context(|| format!("failed to open cert: {cert_path}"))?;
    let key_file = std::fs::File::open(key_path)
        .with_context(|| format!("failed to open key: {key_path}"))?;

    let certs: Vec<_> = rustls_pemfile::certs(&mut BufReader::new(cert_file))
        .collect::<Result<_, _>>()
        .context("failed to parse certificate")?;

    let key = rustls_pemfile::private_key(&mut BufReader::new(key_file))
        .context("failed to parse private key")?
        .ok_or_else(|| anyhow::anyhow!("no private key found in {key_path}"))?;

    // Prefer AES-128-GCM (fastest with AES-NI), pin to TLS 1.3 only.
    let mut provider = aws_lc_rs::default_provider();
    provider.cipher_suites.sort_by_key(|cs| match cs.suite() {
        CipherSuite::TLS13_AES_128_GCM_SHA256 => 0,
        CipherSuite::TLS13_AES_256_GCM_SHA384 => 1,
        _ => 2,
    });

    let config = ServerConfig::builder_with_provider(Arc::new(provider))
        .with_protocol_versions(&[&TLS13])
        .context("TLS 1.3 configuration failed")?
        .with_no_client_auth()
        .with_single_cert(certs, key)
        .context("invalid TLS configuration")?;

    Ok(TlsAcceptor::from(Arc::new(config)))
}
