//! OSTP Client - Stealth tunnel initiator

use std::net::SocketAddr;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

use crate::crypto::{AeadCipher, KeyExchange, PskValidator};
use crate::mimicry::{MimicryEngine, TlsHelloBuilder};
use crate::uot::{encode_frame, decode_frame, FrameFlags};
use bytes::BytesMut;

/// Client configuration
#[derive(Clone)]
pub struct ClientConfig {
    pub server_addr: SocketAddr,
    pub psk: [u8; 32],
    pub country_code: String,
}

impl ClientConfig {
    pub fn new(server: SocketAddr, psk: [u8; 32], country: impl Into<String>) -> Self {
        Self {
            server_addr: server,
            psk,
            country_code: country.into(),
        }
    }
}

/// OSTP Client with TLS mimicry
pub struct OstpClient {
    config: ClientConfig,
    psk_validator: PskValidator,
    mimicry: MimicryEngine,
    cipher: Option<AeadCipher>,
}

impl OstpClient {
    pub fn new(config: ClientConfig) -> Self {
        let psk_validator = PskValidator::new(config.psk);
        Self {
            config,
            psk_validator,
            mimicry: MimicryEngine::new(),
            cipher: None,
        }
    }

    /// Connect to OSTP server with silent handshake
    pub async fn connect(&mut self) -> anyhow::Result<TcpStream> {
        let mut stream = TcpStream::connect(self.config.server_addr).await?;
        tracing::info!("Connected to {}", self.config.server_addr);

        // Phase 1: Generate keypair and send PSK-signed handshake
        let kex = KeyExchange::new();
        let client_pubkey = *kex.public_key();
        
        // Sign the public key with PSK
        let signature = self.psk_validator.sign(client_pubkey.as_bytes());
        
        // Send: [32-byte signature][32-byte pubkey]
        stream.write_all(&signature).await?;
        stream.write_all(client_pubkey.as_bytes()).await?;

        // Phase 2: Receive server's response
        let mut buf = [0u8; 64];
        let n = stream.read_exact(&mut buf).await?;
        if n < 64 {
            anyhow::bail!("Invalid server response");
        }

        let server_sig: [u8; 32] = buf[..32].try_into()?;
        let server_pubkey_bytes: [u8; 32] = buf[32..64].try_into()?;

        // Verify server's signature
        if !self.psk_validator.verify(&server_pubkey_bytes, &server_sig) {
            anyhow::bail!("Server PSK verification failed");
        }

        let server_pubkey = x25519_dalek::PublicKey::from(server_pubkey_bytes);
        let shared_secret = kex.derive_shared(&server_pubkey);
        
        // Derive session key
        let session_key = crate::crypto::derive_session_key(&shared_secret, b"ostp-session-v1");
        self.cipher = Some(AeadCipher::new(&session_key));

        tracing::info!("Session established with server");
        Ok(stream)
    }

    /// Send encrypted data through the tunnel
    pub async fn send(&mut self, stream: &mut TcpStream, data: &[u8]) -> anyhow::Result<()> {
        let cipher = self.cipher.as_mut().ok_or_else(|| anyhow::anyhow!("Not connected"))?;
        
        let encrypted = cipher.encrypt(data).map_err(|_| anyhow::anyhow!("Encryption failed"))?;
        let padding = rand::random::<usize>() % 64; // Random padding 0-63 bytes
        let frame = encode_frame(&encrypted, FrameFlags::DATA, padding)?;
        
        stream.write_all(&frame).await?;
        Ok(())
    }

    /// Receive and decrypt data from the tunnel
    pub async fn recv(&mut self, stream: &mut TcpStream) -> anyhow::Result<Vec<u8>> {
        let cipher = self.cipher.as_ref().ok_or_else(|| anyhow::anyhow!("Not connected"))?;
        
        let mut buf = [0u8; 4096];
        let n = stream.read(&mut buf).await?;
        if n == 0 {
            anyhow::bail!("Connection closed");
        }

        let mut read_buf = BytesMut::from(&buf[..n]);
        
        if let Some((data, _flags)) = decode_frame(&mut read_buf)? {
            let plaintext = cipher.decrypt(&data).map_err(|_| anyhow::anyhow!("Decryption failed"))?;
            Ok(plaintext)
        } else {
            anyhow::bail!("Incomplete frame");
        }
    }

    /// Get SNI for current geo context (for external TLS wrapper if needed)
    pub fn suggested_sni(&self) -> Option<&str> {
        self.mimicry.random_sni(&self.config.country_code)
    }

    /// Build TLS ClientHello for the selected SNI
    pub fn build_tls_hello(&self) -> Option<Vec<u8>> {
        let sni = self.suggested_sni()?;
        Some(TlsHelloBuilder::new(sni).build())
    }
}

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

    #[test]
    fn test_client_config() {
        let addr: SocketAddr = "127.0.0.1:8443".parse().unwrap();
        let psk = [0x42u8; 32];
        let config = ClientConfig::new(addr, psk, "RU");
        assert_eq!(config.country_code, "RU");
    }

    #[test]
    fn test_sni_selection() {
        let addr: SocketAddr = "127.0.0.1:8443".parse().unwrap();
        let config = ClientConfig::new(addr, [0u8; 32], "RU");
        let client = OstpClient::new(config);
        
        let sni = client.suggested_sni();
        assert!(sni.is_some());
        // Should be one of the Russian domains
        let sni = sni.unwrap();
        assert!(sni.ends_with(".ru"));
    }
}
