use anyhow::{ensure, Context, Result};
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, HashSet},
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    time::Duration,
};
use tokio::{
    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
    net::{TcpListener, TcpStream, UdpSocket},
    sync::{mpsc, Mutex, Semaphore},
    task::JoinSet,
    time::Instant,
};
use tokio_util::sync::CancellationToken;

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Protocol {
    #[default]
    Tcp,
    Udp,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct Host {
    pub id: String,
    pub name: String,
    pub enabled: bool,
    pub protocol: Protocol,
    pub listen: SocketAddr,
    pub upstream: SocketAddr,
    /// TCP connect timeout. Kept in the common schema so TCP/UDP rows stay simple.
    pub connect_timeout_seconds: u64,
    /// Inactivity timeout for a TCP relay or a per-client UDP session.
    pub idle_timeout_seconds: u64,
    /// Maximum concurrent TCP connections or UDP client sessions.
    pub max_connections: usize,
}
impl Default for Host {
    fn default() -> Self {
        Self {
            id: String::new(),
            name: String::new(),
            enabled: true,
            protocol: Protocol::Tcp,
            listen: "127.0.0.1:10000".parse().expect("static stream listen"),
            upstream: "127.0.0.1:10001".parse().expect("static stream upstream"),
            connect_timeout_seconds: 5,
            idle_timeout_seconds: 300,
            max_connections: 1024,
        }
    }
}

/// Fixed-cardinality telemetry for one configured L4 stream. The registry is
/// created from the active startup configuration, so stream IDs cannot be
/// attacker-controlled or grow without a configuration/restart boundary.
#[derive(Default, Debug)]
pub struct Metrics {
    pub sessions_total: AtomicU64,
    pub sessions_completed: AtomicU64,
    pub active: AtomicU64,
    pub rejected: AtomicU64,
    pub errors: AtomicU64,
    pub bytes_in: AtomicU64,
    pub bytes_out: AtomicU64,
    pub packets_in: AtomicU64,
    pub packets_out: AtomicU64,
}

#[derive(Debug, Clone, Serialize)]
pub struct MetricsView {
    pub sessions_total: u64,
    pub sessions_completed: u64,
    pub active: u64,
    pub rejected: u64,
    pub errors: u64,
    pub bytes_in: u64,
    pub bytes_out: u64,
    pub packets_in: u64,
    pub packets_out: u64,
}

impl Metrics {
    pub fn view(&self) -> MetricsView {
        let load = |value: &AtomicU64| value.load(Ordering::Relaxed);
        MetricsView {
            sessions_total: load(&self.sessions_total),
            sessions_completed: load(&self.sessions_completed),
            active: load(&self.active),
            rejected: load(&self.rejected),
            errors: load(&self.errors),
            bytes_in: load(&self.bytes_in),
            bytes_out: load(&self.bytes_out),
            packets_in: load(&self.packets_in),
            packets_out: load(&self.packets_out),
        }
    }
}

struct ActiveSession {
    metrics: Arc<Metrics>,
}

impl ActiveSession {
    fn begin(metrics: Arc<Metrics>) -> Self {
        metrics.active.fetch_add(1, Ordering::Relaxed);
        Self { metrics }
    }
}

impl Drop for ActiveSession {
    fn drop(&mut self) {
        self.metrics.active.fetch_sub(1, Ordering::Relaxed);
        self.metrics
            .sessions_completed
            .fetch_add(1, Ordering::Relaxed);
    }
}

pub fn validate(
    hosts: &mut [Host],
    http_listeners: &[crate::vhost::Listener],
    admin: SocketAddr,
    prometheus: Option<SocketAddr>,
) -> Result<()> {
    ensure!(hosts.len() <= 128, "at most 128 stream hosts");
    let mut ids = HashSet::new();
    for host in hosts.iter_mut() {
        ensure!(
            crate::config::valid_id(&host.id) && ids.insert(host.id.clone()),
            "invalid/duplicate stream host id"
        );
        host.name = host.name.trim().to_owned();
        ensure!(
            host.name.len() <= 96 && !host.name.chars().any(char::is_control),
            "stream host name must be at most 96 bytes without controls"
        );
        ensure!(
            host.listen.port() > 0 && host.upstream.port() > 0,
            "stream listener/upstream ports must be nonzero"
        );
        ensure!(
            !host.listen.ip().is_multicast() && !host.upstream.ip().is_multicast(),
            "stream listener/upstream must not use multicast addresses"
        );
        ensure!(
            !host.upstream.ip().is_unspecified(),
            "stream upstream must use a concrete IP address"
        );
        ensure!(
            host.connect_timeout_seconds >= 1 && host.connect_timeout_seconds <= 300,
            "stream connect timeout must be 1..300 seconds"
        );
        ensure!(
            host.idle_timeout_seconds >= 1 && host.idle_timeout_seconds <= 86400,
            "stream idle timeout must be 1..86400 seconds"
        );
        ensure!(
            (1..=65535).contains(&host.max_connections),
            "stream connection/session limit must be 1..65535"
        );
    }

    for (index, host) in hosts.iter().enumerate() {
        if !host.enabled {
            continue;
        }
        for other in hosts.iter().skip(index + 1).filter(|item| item.enabled) {
            if host.protocol == other.protocol {
                ensure!(
                    !crate::vhost::socket_overlap(host.listen, other.listen),
                    "stream listeners {} and {} overlap",
                    host.id,
                    other.id
                );
            }
        }
        match host.protocol {
            Protocol::Tcp => {
                ensure!(
                    !crate::vhost::socket_overlap(host.listen, admin),
                    "TCP stream {} overlaps administration listener",
                    host.id
                );
                if let Some(metrics) = prometheus {
                    ensure!(
                        !crate::vhost::socket_overlap(host.listen, metrics),
                        "TCP stream {} overlaps Prometheus listener",
                        host.id
                    );
                }
                for listener in http_listeners {
                    ensure!(
                        !crate::vhost::socket_overlap(host.listen, listener.address),
                        "TCP stream {} overlaps HTTP/TLS listener",
                        host.id
                    );
                }
            }
            Protocol::Udp => {
                for listener in http_listeners
                    .iter()
                    .filter(|listener| listener.http3.is_some())
                {
                    ensure!(
                        !crate::vhost::socket_overlap(host.listen, listener.address),
                        "UDP stream {} overlaps HTTP/3 listener",
                        host.id
                    );
                }
            }
        }
    }

    // Reject direct loops into another RustProxy listener of the same transport.
    for host in hosts.iter().filter(|item| item.enabled) {
        match host.protocol {
            Protocol::Tcp => {
                ensure!(
                    !crate::vhost::socket_overlap(host.upstream, admin)
                        && prometheus
                            .is_none_or(|addr| !crate::vhost::socket_overlap(host.upstream, addr))
                        && !http_listeners
                            .iter()
                            .any(|listener| crate::vhost::socket_overlap(
                                host.upstream,
                                listener.address
                            ))
                        && !hosts
                            .iter()
                            .filter(|other| other.enabled && other.protocol == Protocol::Tcp)
                            .any(|other| crate::vhost::socket_overlap(host.upstream, other.listen)),
                    "TCP stream {} upstream points at a RustProxy listener",
                    host.id
                );
            }
            Protocol::Udp => {
                ensure!(
                    !http_listeners
                        .iter()
                        .filter(|listener| listener.http3.is_some())
                        .any(|listener| crate::vhost::socket_overlap(
                            host.upstream,
                            listener.address
                        ))
                        && !hosts
                            .iter()
                            .filter(|other| other.enabled && other.protocol == Protocol::Udp)
                            .any(|other| crate::vhost::socket_overlap(host.upstream, other.listen)),
                    "UDP stream {} upstream points at a RustProxy UDP listener",
                    host.id
                );
            }
        }
    }
    Ok(())
}

pub async fn serve_tcp(
    host: Host,
    listener: TcpListener,
    shutdown: CancellationToken,
    grace: Duration,
    metrics: Arc<Metrics>,
) -> Result<()> {
    let slots = Arc::new(Semaphore::new(host.max_connections));
    let mut tasks = JoinSet::new();
    loop {
        tokio::select! {
            biased;
            _ = shutdown.cancelled() => break,
            _ = tasks.join_next(), if !tasks.is_empty() => {},
            accepted = listener.accept() => {
                let (client, peer) = accepted.context("stream TCP accept failed")?;
                let Ok(permit) = slots.clone().try_acquire_owned() else {
                    metrics.rejected.fetch_add(1, Ordering::Relaxed);
                    drop(client);
                    continue;
                };
                metrics.sessions_total.fetch_add(1, Ordering::Relaxed);
                let cfg = host.clone();
                let stop = shutdown.clone();
                let stream_metrics = metrics.clone();
                let active = ActiveSession::begin(stream_metrics.clone());
                tasks.spawn(async move {
                    let _active = active;
                    let _permit = permit;
                    let upstream = match tokio::time::timeout(
                        Duration::from_secs(cfg.connect_timeout_seconds),
                        TcpStream::connect(cfg.upstream),
                    ).await {
                        Ok(Ok(stream)) => stream,
                        Ok(Err(error)) => {
                            stream_metrics.errors.fetch_add(1, Ordering::Relaxed);
                            tracing::debug!(stream=%cfg.id, %peer, %error, "TCP stream upstream connect failed");
                            return;
                        },
                        Err(_) => {
                            stream_metrics.errors.fetch_add(1, Ordering::Relaxed);
                            tracing::debug!(stream=%cfg.id, %peer, "TCP stream upstream connect timed out");
                            return;
                        },
                    };
                    if let Err(error) = relay_tcp(client, upstream, stop, Duration::from_secs(cfg.idle_timeout_seconds), &stream_metrics).await {
                        stream_metrics.errors.fetch_add(1, Ordering::Relaxed);
                        tracing::debug!(stream=%cfg.id, %peer, %error, "TCP stream relay ended");
                    }
                });
            }
        }
    }
    let _ = tokio::time::timeout(grace, async { while tasks.join_next().await.is_some() {} }).await;
    tasks.abort_all();
    while tasks.join_next().await.is_some() {}
    Ok(())
}

async fn copy_with_activity<R, W>(
    mut reader: R,
    mut writer: W,
    activity: tokio::sync::watch::Sender<Instant>,
    transferred: &AtomicU64,
) -> std::io::Result<u64>
where
    R: AsyncRead + Unpin,
    W: AsyncWrite + Unpin,
{
    let mut total = 0u64;
    let mut buffer = [0u8; 32 * 1024];
    loop {
        let read = reader.read(&mut buffer).await?;
        if read == 0 {
            writer.shutdown().await?;
            return Ok(total);
        }
        writer.write_all(&buffer[..read]).await?;
        total = total.saturating_add(read as u64);
        transferred.fetch_add(read as u64, Ordering::Relaxed);
        activity.send_replace(Instant::now());
    }
}

async fn relay_tcp(
    client: TcpStream,
    upstream: TcpStream,
    shutdown: CancellationToken,
    idle: Duration,
    metrics: &Metrics,
) -> Result<()> {
    let (client_read, client_write) = client.into_split();
    let (upstream_read, upstream_write) = upstream.into_split();
    let (activity, mut activity_rx) = tokio::sync::watch::channel(Instant::now());
    let c2u = copy_with_activity(
        client_read,
        upstream_write,
        activity.clone(),
        &metrics.bytes_in,
    );
    let u2c = copy_with_activity(
        upstream_read,
        client_write,
        activity.clone(),
        &metrics.bytes_out,
    );
    drop(activity);
    tokio::pin!(c2u);
    tokio::pin!(u2c);
    let mut c2u_done = false;
    let mut u2c_done = false;
    let deadline = tokio::time::sleep_until(Instant::now() + idle);
    tokio::pin!(deadline);
    loop {
        tokio::select! {
            _ = shutdown.cancelled() => break,
            result = &mut c2u, if !c2u_done => { result?; c2u_done = true; if u2c_done { break; } },
            result = &mut u2c, if !u2c_done => { result?; u2c_done = true; if c2u_done { break; } },
            changed = activity_rx.changed() => {
                if changed.is_err() { break; }
                deadline.as_mut().reset(*activity_rx.borrow() + idle);
            },
            _ = &mut deadline => break,
        }
    }
    Ok(())
}

#[derive(Clone)]
struct UdpSession {
    id: u64,
    tx: mpsc::Sender<Bytes>,
}