//! OSTP Server - Stealth relay endpoint

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::RwLock;

use crate::crypto::{AeadCipher, KeyExchange, PskValidator};
use crate::uot::decode_frame;
use bytes::BytesMut;

/// Server configuration
#[derive(Clone)]
pub struct ServerConfig {
    pub listen_addr: SocketAddr,
    pub psk: [u8; 32],
    pub max_connections: usize,
}

impl ServerConfig {
    pub fn new(addr: SocketAddr, psk: [u8; 32]) -> Self {
        Self {
            listen_addr: addr,
            psk,
            max_connections: 1024,
        }
    }
}

/// Active connection state
#[allow(dead_code)]
struct Connection {
    cipher: AeadCipher,
    user_id: Option<uuid::Uuid>,
    bytes_rx: u64,
    bytes_tx: u64,
}

/// OSTP Server with silent handshake
pub struct OstpServer {
    config: ServerConfig,
    #[allow(dead_code)]
    psk_validator: PskValidator,
    connections: Arc<RwLock<HashMap<SocketAddr, Connection>>>,
}

impl OstpServer {
    pub fn new(config: ServerConfig) -> Self {
        let psk_validator = PskValidator::new(config.psk);
        Self {
            config,
            psk_validator,
            connections: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Start the server (main async loop)
    pub async fn run(&self) -> anyhow::Result<()> {
        let listener = TcpListener::bind(self.config.listen_addr).await?;
        tracing::info!("OSTP server listening on {}", self.config.listen_addr);

        loop {
            match listener.accept().await {
                Ok((stream, addr)) => {
                    let psk_validator = PskValidator::new(self.config.psk);
                    let connections = Arc::clone(&self.connections);
                    
                    tokio::spawn(async move {
                        if let Err(e) = Self::handle_connection(stream, addr, psk_validator, connections).await {
                            tracing::debug!("Connection {} closed: {}", addr, e);
                        }
                    });
                }
                Err(e) => {
                    tracing::warn!("Accept error: {}", e);
                }
            }
        }
    }

    async fn handle_connection(
        mut stream: TcpStream,
        addr: SocketAddr,
        psk_validator: PskValidator,
        connections: Arc<RwLock<HashMap<SocketAddr, Connection>>>,
    ) -> anyhow::Result<()> {
        let mut buf = [0u8; 4096];
        
        // Phase 1: Silent PSK handshake (32-byte signature must be first)
        let n = stream.read(&mut buf).await?;
        if n < 64 {
            // Silent drop - don't reveal port is open
            return Ok(());
        }

        let signature: [u8; 32] = buf[..32].try_into()?;
        let payload = &buf[32..n];

        if !psk_validator.verify(payload, &signature) {
            // Silent drop on invalid PSK - core anti-probing defense
            tracing::trace!("Silent drop for {}: invalid PSK", addr);
            return Ok(());
        }

        // Phase 2: X25519 key exchange
        // Payload contains client's public key (32 bytes)
        if payload.len() < 32 {
            return Ok(());
        }
        
        let client_pubkey_bytes: [u8; 32] = payload[..32].try_into()?;
        let client_pubkey = x25519_dalek::PublicKey::from(client_pubkey_bytes);
        
        let server_kex = KeyExchange::new();
        let server_pubkey = *server_kex.public_key();
        let shared_secret = server_kex.derive_shared(&client_pubkey);
        
        // Derive session key
        let session_key = crate::crypto::derive_session_key(&shared_secret, b"ostp-session-v1");
        let cipher = AeadCipher::new(&session_key);

        // Send server's public key (signed with PSK)
        let server_response = server_pubkey.as_bytes();
        let response_sig = psk_validator.sign(server_response);
        stream.write_all(&response_sig).await?;
        stream.write_all(server_response).await?;

        // Store connection
        {
            let mut conns = connections.write().await;
            conns.insert(addr, Connection {
                cipher,
                user_id: None,
                bytes_rx: 0,
                bytes_tx: 0,
            });
        }

        tracing::info!("Session established with {}", addr);

        // Phase 3: Encrypted relay loop
        let mut read_buf = BytesMut::with_capacity(65536);
        
        loop {
            let n = stream.read(&mut buf).await?;
            if n == 0 {
                break;
            }

            read_buf.extend_from_slice(&buf[..n]);

            // Process frames
            while let Some((data, flags)) = decode_frame(&mut read_buf)? {
                let conns = connections.read().await;
                if let Some(conn) = conns.get(&addr) {
                    match conn.cipher.decrypt(&data) {
                        Ok(plaintext) => {
                            if flags.is_data() {
                                // TODO: Route decrypted data to destination
                                tracing::trace!("Received {} bytes from {}", plaintext.len(), addr);
                            }
                        }
                        Err(_) => {
                            tracing::warn!("Decrypt failed for {}", addr);
                            break;
                        }
                    }
                }
            }
        }

        // Cleanup
        connections.write().await.remove(&addr);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_server_config() {
        let addr: SocketAddr = "127.0.0.1:8443".parse().unwrap();
        let psk = [0x42u8; 32];
        let config = ServerConfig::new(addr, psk);
        assert_eq!(config.listen_addr, addr);
        assert_eq!(config.max_connections, 1024);
    }
}
