//! Authenticated, opaque-datagram direct UDP transport.

use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::{
        atomic::{AtomicU32, Ordering},
        Arc,
    },
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use anyhow::{anyhow, Context, Result};
use bytes::Bytes;
use rand::{rngs::OsRng, RngCore};
use tokio::{
    net::{lookup_host, UdpSocket},
    sync::Mutex,
    task::JoinSet,
    time::{self, timeout},
};
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};

use crate::{
    config::{ClientConfig, ServerConfig},
    protocol::{
        self, auth,
        direct::{
            AcceptPacket, DataPacket, HelloPacket, Packet, PacketType, RejectPacket, RejectReason,
            SessionPacket,
        },
        replay::{NonceReplayCache, ReplayWindow},
    },
    stats::Stats,
};

use super::{TransportConnector, TransportError, TransportFuture, TransportSession};

const DEFAULT_MAX_DATAGRAM: usize = 1200;
const MAX_UDP_DATAGRAM: usize = u16::MAX as usize;

fn unix_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

async fn resolve_server(server: &str, port: u16) -> Result<SocketAddr> {
    if let Ok(addr) = server.parse::<SocketAddr>() {
        return Ok(addr);
    }
    let endpoint = format!("{server}:{port}");
    lookup_host(endpoint.clone())
        .await?
        .next()
        .ok_or_else(|| anyhow!("could not resolve {endpoint}"))
}

fn direct_bind_addr(listen: &str, port: u16) -> Result<String> {
    if listen.parse::<SocketAddr>().is_ok() {
        Ok(listen.to_owned())
    } else {
        Ok(format!("{listen}:{port}"))
    }
}

/// Established client-side direct UDP session.
pub struct DirectUdpSession {
    socket: Arc<UdpSocket>,
    peer: Mutex<SocketAddr>,
    session_id: u64,
    session_key: [u8; auth::SESSION_KEY_LEN],
    send_sequence: AtomicU32,
    receive_window: Mutex<ReplayWindow>,
    max_datagram: usize,
    closed: CancellationToken,
}

impl DirectUdpSession {
    pub fn session_id(&self) -> u64 {
        self.session_id
    }
    pub fn peer(&self) -> TransportFuture<'_, SocketAddr> {
        Box::pin(async move { Ok(*self.peer.lock().await) })
    }

    pub async fn connect(config: &ClientConfig) -> Result<Self, TransportError> {
        let secret = config.client.secret.load_bytes("client").map_err(|e| {
            TransportError::Io(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                e.to_string(),
            ))
        })?;
        let peer = resolve_server(&config.client.server, config.transport.direct_udp.port)
            .await
            .map_err(|e| {
                TransportError::Io(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    e.to_string(),
                ))
            })?;
        let socket = UdpSocket::bind("0.0.0.0:0").await?;
        let socket = Arc::new(socket);
        let mut nonce = [0u8; 16];
        OsRng.fill_bytes(&mut nonce);
        let hello = HelloPacket::authenticated(
            &secret,
            Bytes::copy_from_slice(config.client.id.as_bytes()),
            unix_timestamp(),
            nonce,
        )?;
        let encoded = hello.encode()?;
        let handshake_timeout =
            Duration::from_millis(config.transport.auto.probe_timeout_ms.max(1))
                .min(Duration::from_secs(30));
        for _attempt in 0..config.transport.auto.retries.max(1) {
            socket.send_to(&encoded, peer).await?;
            let mut buf = vec![0u8; MAX_UDP_DATAGRAM];
            match timeout(handshake_timeout, socket.recv_from(&mut buf)).await {
                Ok(Ok((len, source))) => {
                    let packet = match protocol::parse_packet(&buf[..len]) {
                        Ok(v) => v,
                        Err(e) => {
                            debug!(error=%e, "ignoring malformed handshake response");
                            continue;
                        }
                    };
                    match packet {
                        Packet::Accept(accept) => {
                            if source.ip() != peer.ip()
                                || !accept.verify_authenticator(&secret, &nonce)
                            {
                                continue;
                            }
                            let max_datagram = (config.network.max_datagram_size as usize)
                                .min(accept.max_datagram as usize);
                            let session_key = auth::derive_session_key(
                                &secret,
                                &nonce,
                                &accept.server_nonce,
                                accept.session_id,
                            );
                            info!(session_id=accept.session_id, remote_addr=%source, max_datagram, "direct UDP session authenticated");
                            return Ok(Self {
                                socket,
                                peer: Mutex::new(source),
                                session_id: accept.session_id,
                                session_key,
                                send_sequence: AtomicU32::new(auth::generate_sequence_number()),
                                receive_window: Mutex::new(ReplayWindow::new()),
                                max_datagram,
                                closed: CancellationToken::new(),
                            });
                        }
                        Packet::Reject(reject) => {
                            return Err(TransportError::AuthenticationRejected(reject.reason))
                        }
                        _ => continue,
                    }
                }
                Ok(Err(e)) => return Err(e.into()),
                Err(_) => continue,
            }
        }
        Err(TransportError::HandshakeTimeout)
    }

    async fn send_packet(&self, bytes: &[u8]) -> Result<(), TransportError> {
        if self.closed.is_cancelled() {
            return Err(TransportError::Closed);
        }
        let peer = *self.peer.lock().await;
        self.socket.send_to(bytes, peer).await?;
        Ok(())
    }

    /// Send an authenticated control keepalive.  It intentionally shares the
    /// outgoing DATA sequence space only within this direction.
    pub async fn send_keepalive(&self) -> Result<(), TransportError> {
        let sequence = self.send_sequence.fetch_add(1, Ordering::Relaxed);
        let packet = SessionPacket::authenticated(
            PacketType::Keepalive,
            &self.session_key,
            self.session_id,
            sequence,
        );
        self.send_packet(&packet.encode()?).await
    }

    async fn receive_valid(&self) -> Result<(Packet, SocketAddr), TransportError> {
        let mut buf = vec![0u8; MAX_UDP_DATAGRAM];
        loop {
            let (len, source) = self.socket.recv_from(&mut buf).await?;
            let packet = match protocol::parse_packet(&buf[..len]) {
                Ok(p) => p,
                Err(e) => {
                    debug!(error=%e, "dropping malformed direct UDP packet");
                    continue;
                }
            };
            let valid = match &packet {
                Packet::Data(p) => {
                    p.session_id == self.session_id && p.verify_authenticator(&self.session_key)
                }
                Packet::Keepalive(p) | Packet::Close(p) => {
                    p.session_id == self.session_id && p.verify_authenticator(&self.session_key)
                }
                _ => false,
            };
            if !valid {
                continue;
            }
            let sequence = match &packet {
                Packet::Data(p) => p.sequence,
                Packet::Keepalive(p) | Packet::Close(p) => p.sequence,
                _ => unreachable!(),
            };
            if !self.receive_window.lock().await.accept(sequence) {
                continue;
            }
            *self.peer.lock().await = source;
            return Ok((packet, source));
        }
    }
}

impl TransportSession for DirectUdpSession {
    fn send_datagram<'a>(&'a self, data: Bytes) -> TransportFuture<'a, ()> {
        Box::pin(async move {
            if data.len() > self.max_datagram {
                return Err(TransportError::Oversized {
                    maximum: self.max_datagram,
                });
            }
            let sequence = self.send_sequence.fetch_add(1, Ordering::Relaxed);
            let packet =
                DataPacket::authenticated(&self.session_key, self.session_id, sequence, data);
            self.send_packet(&packet.encode()?).await
        })
    }
    fn recv_datagram<'a>(&'a self) -> TransportFuture<'a, Bytes> {
        Box::pin(async move {
            loop {
                match self.receive_valid().await? {
                    (Packet::Data(packet), _) => return Ok(packet.payload),
                    (Packet::Keepalive(_), _) => continue,
                    (Packet::Close(_), _) => {
                        self.closed.cancel();
                        return Err(TransportError::Closed);
                    }
                    _ => continue,
                }
            }
        })
    }
    fn close<'a>(&'a self) -> TransportFuture<'a, ()> {
        Box::pin(async move {
            if self.closed.is_cancelled() {
                return Ok(());
            }
            let sequence = self.send_sequence.fetch_add(1, Ordering::Relaxed);
            let packet = SessionPacket::authenticated(
                PacketType::Close,
                &self.session_key,
                self.session_id,
                sequence,
            );
            let _ = self.send_packet(&packet.encode()?).await;
            self.closed.cancel();
            Ok(())
        })
    }
    fn max_datagram_size(&self) -> usize {
        self.max_datagram
    }
}

pub struct DirectUdpConnector {
    config: ClientConfig,
}
impl DirectUdpConnector {
    pub fn new(config: ClientConfig) -> Self {
        Self { config }
    }
}
impl TransportConnector for DirectUdpConnector {
    fn connect<'a>(&'a self) -> TransportFuture<'a, Box<dyn TransportSession>> {
        Box::pin(async move {
            Ok(Box::new(DirectUdpSession::connect(&self.config).await?)
                as Box<dyn TransportSession>)
        })
    }
}

struct ServerSession {
    id: u64,
    client_id: String,
    key: [u8; auth::SESSION_KEY_LEN],
    peer: Mutex<SocketAddr>,
    backend: Arc<UdpSocket>,
    receive_window: Mutex<ReplayWindow>,
    send_sequence: AtomicU32,
    last_activity: Mutex<Instant>,
    max_datagram: usize,
    cancel: CancellationToken,
}

/// Multi-client direct UDP server. Each authenticated session gets its own
/// connected backend socket and therefore natural response demultiplexing.
pub struct DirectUdpServer {
    socket: Arc<UdpSocket>,
    backend_target: SocketAddr,
    clients: HashMap<String, Vec<u8>>,
    sessions: Arc<Mutex<HashMap<u64, Arc<ServerSession>>>>,
    nonces: Arc<Mutex<NonceReplayCache>>,
    pending: Arc<Mutex<HashMap<SocketAddr, Instant>>>,
    config: ServerConfig,
    stats: Arc<Stats>,
}

impl DirectUdpServer {
    pub async fn bind(config: ServerConfig) -> Result<Self> {
        let bind = direct_bind_addr(&config.server.listen, config.transports.direct_udp.port)?;
        let socket = Arc::new(
            UdpSocket::bind(&bind)
                .await
                .context("bind direct UDP listener")?,
        );
        let backend_target = config
            .backend
            .target
            .parse::<SocketAddr>()
            .context("invalid backend target")?;
        let mut clients = HashMap::new();
        for client in &config.auth.clients {
            clients.insert(
                client.id.clone(),
                client
                    .secret
                    .load_bytes(&format!("auth client {}", client.id))?,
            );
        }
        Ok(Self {
            socket,
            backend_target,
            clients,
            sessions: Arc::new(Mutex::new(HashMap::new())),
            nonces: Arc::new(Mutex::new(NonceReplayCache::new(Duration::from_secs(120)))),
            pending: Arc::new(Mutex::new(HashMap::new())),
            config,
            stats: Arc::new(Stats::default()),
        })
    }
    pub fn local_addr(&self) -> Result<SocketAddr> {
        Ok(self.socket.local_addr()?)
    }
    pub fn stats(&self) -> Arc<Stats> {
        self.stats.clone()
    }
    pub async fn active_sessions(&self) -> usize {
        self.sessions.lock().await.len()
    }

    async fn send_reject(&self, peer: SocketAddr, reason: RejectReason) {
        if let Ok(bytes) = RejectPacket::new(reason).encode() {
            if self.socket.send_to(&bytes, peer).await.is_ok() {
                self.stats.add_sent(bytes.len());
            }
        }
        self.stats
            .authentication_rejections
            .fetch_add(1, Ordering::Relaxed);
        self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
    }

    async fn handle_hello(
        &self,
        hello: protocol::HelloPacket,
        peer: SocketAddr,
    ) -> Result<Option<Arc<ServerSession>>> {
        self.stats
            .handshake_attempts
            .fetch_add(1, Ordering::Relaxed);
        let now = Instant::now();
        {
            let mut pending = self.pending.lock().await;
            let expiry = Duration::from_millis(self.config.limits.handshake_timeout_ms.max(1));
            pending.retain(|_, seen| now.duration_since(*seen) < expiry);
            if !pending.contains_key(&peer)
                && pending.len() >= self.config.limits.max_pending_handshakes
            {
                self.stats
                    .rate_limited_attempts
                    .fetch_add(1, Ordering::Relaxed);
                drop(pending);
                self.send_reject(peer, RejectReason::RateLimited).await;
                return Ok(None);
            }
            pending.insert(peer, now);
        }
        let source_ip = peer.ip();
        let active = self.sessions.lock().await;
        let sessions_for_ip = {
            let mut count = 0;
            for session in active.values() {
                if session.peer.lock().await.ip() == source_ip {
                    count += 1;
                }
            }
            count
        };
        if active.len() >= self.config.limits.max_sessions
            || sessions_for_ip >= self.config.limits.max_sessions_per_ip
        {
            drop(active);
            self.pending.lock().await.remove(&peer);
            self.send_reject(peer, RejectReason::ServerBusy).await;
            return Ok(None);
        }
        drop(active);
        let Some(secret) = self
            .clients
            .get(std::str::from_utf8(&hello.client_id).unwrap_or(""))
        else {
            self.pending.lock().await.remove(&peer);
            self.send_reject(peer, RejectReason::AuthenticationFailed)
                .await;
            return Ok(None);
        };
        if !hello.is_fresh(SystemTime::now(), Duration::from_secs(60))
            || !hello.verify_authenticator(secret)
        {
            self.stats
                .authentication_failures
                .fetch_add(1, Ordering::Relaxed);
            self.pending.lock().await.remove(&peer);
            self.send_reject(peer, RejectReason::AuthenticationFailed)
                .await;
            return Ok(None);
        }
        if !self.nonces.lock().await.check_and_record(
            &hello.client_id,
            hello.client_nonce,
            Instant::now(),
        ) {
            self.pending.lock().await.remove(&peer);
            self.send_reject(peer, RejectReason::ReplayDetected).await;
            return Ok(None);
        }
        let backend = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
        backend.connect(self.backend_target).await?;
        let mut session_id = auth::generate_session_id();
        let mut sessions = self.sessions.lock().await;
        while sessions.contains_key(&session_id) {
            session_id = auth::generate_session_id();
        }
        let mut server_nonce = [0u8; 16];
        OsRng.fill_bytes(&mut server_nonce);
        let idle = self.config.session.idle_timeout_secs.min(u32::MAX as u64) as u32;
        let max_datagram = DEFAULT_MAX_DATAGRAM as u16;
        let key = auth::derive_session_key(secret, &hello.client_nonce, &server_nonce, session_id);
        let session = Arc::new(ServerSession {
            id: session_id,
            client_id: String::from_utf8_lossy(&hello.client_id).into_owned(),
            key,
            peer: Mutex::new(peer),
            backend,
            receive_window: Mutex::new(ReplayWindow::new()),
            send_sequence: AtomicU32::new(auth::generate_sequence_number()),
            last_activity: Mutex::new(Instant::now()),
            max_datagram: max_datagram as usize,
            cancel: CancellationToken::new(),
        });
        sessions.insert(session_id, session.clone());
        drop(sessions);
        self.pending.lock().await.remove(&peer);
        let accept = AcceptPacket::authenticated(
            secret,
            session_id,
            hello.client_nonce,
            server_nonce,
            idle,
            max_datagram,
        );
        let accept_bytes = accept.encode()?;
        self.socket.send_to(&accept_bytes, peer).await?;
        self.stats.add_sent(accept_bytes.len());
        info!(session_id, client_id=%session.client_id, remote_addr=%peer, "direct UDP client authenticated");
        Ok(Some(session))
    }

    async fn remove_session(&self, id: u64, expired: bool) {
        if let Some(session) = self.sessions.lock().await.remove(&id) {
            session.cancel.cancel();
            if expired {
                self.stats.expired_sessions.fetch_add(1, Ordering::Relaxed);
            }
        }
    }

    pub async fn run(self) -> Result<()> {
        let mut tasks = JoinSet::new();
        let mut buf = vec![0u8; MAX_UDP_DATAGRAM];
        let mut cleanup = time::interval(Duration::from_secs(1));
        let signal = shutdown_signal();
        tokio::pin!(signal);
        loop {
            tokio::select! {
                _ = &mut signal => { info!("direct UDP server shutting down"); break; }
                Some(done) = tasks.join_next() => { if let Ok(id) = done { self.remove_session(id, false).await; } }
                _ = cleanup.tick() => {
                    let now = Instant::now();
                    let expired: Vec<u64> = self.sessions.lock().await.iter().filter_map(|(&id,s)| async_ready_expired(s, now, Duration::from_secs(self.config.session.idle_timeout_secs)).map(|_| id)).collect();
                    for id in expired { self.remove_session(id, true).await; }
                }
                received = self.socket.recv_from(&mut buf) => {
                    let (len, peer) = match received { Ok(v) => v, Err(e) => { warn!(error=%e, "direct UDP receive failed"); continue; } };
                    self.stats.add_received(len);
                    let packet = match protocol::parse_packet(&buf[..len]) { Ok(p) => p, Err(e) => { self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed); debug!(remote_addr=%peer,error=%e,"dropping malformed direct UDP packet"); continue; } };
                    match packet {
                        Packet::Hello(hello) => match self.handle_hello(hello, peer).await {
                            Ok(Some(session)) => {
                                let listener = self.socket.clone(); let backend = session.backend.clone(); let session_clone = session.clone(); let stats = self.stats.clone();
                                tasks.spawn(async move { backend_loop(listener, backend, session_clone, stats).await });
                            }
                            Ok(None) => {}
                            Err(error) => warn!(remote_addr=%peer, error=%error, "failed to establish direct UDP session"),
                        },
                        Packet::Data(data) => if let Err(error) = self.handle_session_packet(PacketType::Data, data.session_id, data.sequence, data.auth_tag, data.payload, peer).await { self.stats.backend_errors.fetch_add(1, Ordering::Relaxed); self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed); warn!(session_id=data.session_id, error=%error, "backend forwarding failed"); },
                        Packet::Keepalive(p) => if let Err(error) = self.handle_session_packet(p.packet_type, p.session_id, p.sequence, p.auth_tag, Bytes::new(), peer).await { warn!(session_id=p.session_id, error=%error, "keepalive handling failed"); },
                        Packet::Close(p) => { if let Err(error) = self.handle_session_packet(p.packet_type, p.session_id, p.sequence, p.auth_tag, Bytes::new(), peer).await { warn!(session_id=p.session_id, error=%error, "close handling failed"); } self.remove_session(p.session_id, false).await; },
                        _ => {}
                    }
                }
            }
        }
        let ids: Vec<u64> = self.sessions.lock().await.keys().copied().collect();
        for id in ids {
            self.remove_session(id, false).await;
        }
        while tasks.join_next().await.is_some() {}
        Ok(())
    }

    async fn handle_session_packet(
        &self,
        packet_type: PacketType,
        id: u64,
        sequence: u32,
        tag: [u8; 16],
        payload: Bytes,
        peer: SocketAddr,
    ) -> Result<()> {
        let session = match self.sessions.lock().await.get(&id).cloned() {
            Some(s) => s,
            None => {
                self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
                return Ok(());
            }
        };
        if payload.len() > session.max_datagram {
            self.stats
                .oversized_datagrams
                .fetch_add(1, Ordering::Relaxed);
            self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
            return Ok(());
        }
        if !auth::verify_data_auth_tag(
            &session.key,
            protocol::PROTOCOL_VERSION,
            packet_type,
            0,
            id,
            sequence,
            &payload,
            &tag,
        ) {
            self.stats
                .authentication_failures
                .fetch_add(1, Ordering::Relaxed);
            self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
            return Ok(());
        }
        if !session.receive_window.lock().await.accept(sequence) {
            self.stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
            return Ok(());
        }
        *session.peer.lock().await = peer;
        *session.last_activity.lock().await = Instant::now();
        match packet_type {
            PacketType::Data => {
                session.backend.send(&payload).await?;
            }
            PacketType::Close => session.cancel.cancel(),
            PacketType::Keepalive => {}
            _ => {}
        }
        Ok(())
    }
}

async fn backend_loop(
    listener: Arc<UdpSocket>,
    backend: Arc<UdpSocket>,
    session: Arc<ServerSession>,
    stats: Arc<Stats>,
) -> u64 {
    let mut buf = vec![0u8; MAX_UDP_DATAGRAM];
    loop {
        tokio::select! {
            _ = session.cancel.cancelled() => break,
            received = backend.recv(&mut buf) => {
                let len = match received { Ok(len) => len, Err(_) => break };
                if len > session.max_datagram {
                    stats.oversized_datagrams.fetch_add(1, Ordering::Relaxed);
                    stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
                    continue;
                }
                let sequence = session.send_sequence.fetch_add(1, Ordering::Relaxed);
                let packet = DataPacket::authenticated(&session.key, session.id, sequence, Bytes::copy_from_slice(&buf[..len]));
                let bytes = match packet.encode() { Ok(v) => v, Err(_) => continue };
                let peer = *session.peer.lock().await;
                if listener.send_to(&bytes, peer).await.is_err() {
                    stats.backend_errors.fetch_add(1, Ordering::Relaxed);
                    stats.datagrams_dropped.fetch_add(1, Ordering::Relaxed);
                    break;
                }
                stats.add_sent(bytes.len());
            }
        }
    }
    session.id
}

fn async_ready_expired(
    session: &Arc<ServerSession>,
    now: Instant,
    timeout: Duration,
) -> Option<u64> {
    session.last_activity.try_lock().ok().and_then(|last| {
        if now.duration_since(*last) >= timeout {
            Some(session.id)
        } else {
            None
        }
    })
}

async fn shutdown_signal() {
    #[cfg(unix)]
    {
        use tokio::signal::unix::{signal, SignalKind};
        let mut term = signal(SignalKind::terminate()).expect("install SIGTERM handler");
        tokio::select! { _ = tokio::signal::ctrl_c() => {}, _ = term.recv() => {} }
    }
    #[cfg(not(unix))]
    {
        let _ = tokio::signal::ctrl_c().await;
    }
}