// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 LLMHosts.com Contributors
//
//! Noise NK handshake and transport layer for LLMHosts relay.
//!
//! # Protocol
//!
//! Uses `Noise_NK_25519_AESGCM_SHA256` — the NK (No initiator key, Known
//! responder key) pattern:
//!
//! ```text
//! <- s          (server static pubkey pre-shared to client)
//! ...
//! -> e, es      (client ephemeral + DH with server static)
//! <- e, ee      (server ephemeral + DH ephemeral-ephemeral)
//! ```
//!
//! After handshake, all WebSocket frames are encrypted with AESGCM-256.
//! The relay server (VPS) cannot read request/response payloads — it only
//! routes opaque encrypted binary frames.
//!
//! # Usage
//!
//! ```rust,ignore
//! let session = NoiseSession::handshake_initiator(&server_pubkey_hex, &mut ws_stream).await?;
//! let ciphertext = session.encrypt(b"hello relay")?;
//! let plaintext = session.decrypt(&ciphertext)?;
//! ```

use std::sync::Mutex;

/// Noise NK transport session — wraps snow TransportState behind a Mutex.
///
/// `Mutex` is required because TransportState tracks nonces for both send
/// and receive. M3 will upgrade to per-direction stateless transport.
pub struct NoiseSession {
    transport: Mutex<snow::TransportState>,
}

/// Parameters for the Noise NK handshake.
const NOISE_PARAMS: &str = "Noise_NK_25519_AESGCM_SHA256";

/// Maximum plaintext size per Noise message (65535 - 16 bytes AEAD tag).
pub const NOISE_MAX_MSG: usize = 65519;

impl NoiseSession {
    /// Wrap an existing `snow::TransportState` (post-handshake) in a session.
    ///
    /// Used by `lib.rs` which performs the handshake inline over the WS split.
    pub fn from_transport(transport: snow::TransportState) -> Self {
        Self { transport: Mutex::new(transport) }
    }

    /// Perform the Noise NK handshake as the **initiator** (device side).
    ///
    /// `server_pubkey_hex` — hex-encoded 32-byte Curve25519 public key of the
    /// relay server. Shipped with the LLMHosts package or fetched during auth.
    ///
    /// `ws` — reference to the mutable write half of the WebSocket, used to
    /// exchange the two handshake messages (→ and ←).
    pub async fn handshake_initiator<S>(
        server_pubkey_hex: &str,
        ws: &mut S,
    ) -> Result<Self, NoiseError>
    where
        S: NoiseTransport,
    {
        let server_pubkey = decode_pubkey(server_pubkey_hex)?;

        let params: snow::params::NoiseParams = NOISE_PARAMS
            .parse()
            .map_err(|e: snow::Error| NoiseError::Init(e.to_string()))?;

        let mut initiator = snow::Builder::new(params)
            .remote_public_key(&server_pubkey)
            .build_initiator()
            .map_err(|e| NoiseError::Init(e.to_string()))?;

        // Message 1: -> e, es
        let mut buf = vec![0u8; 256];
        let len = initiator
            .write_message(&[], &mut buf)
            .map_err(|e| NoiseError::Handshake(e.to_string()))?;
        ws.send_raw(&buf[..len]).await?;

        // Message 2: <- e, ee
        let msg2 = ws.recv_raw().await?;
        initiator
            .read_message(&msg2, &mut buf)
            .map_err(|e| NoiseError::Handshake(e.to_string()))?;

        // Handshake complete — transition to transport mode
        let transport = initiator
            .into_transport_mode()
            .map_err(|e| NoiseError::Handshake(e.to_string()))?;

        tracing::info!("Noise NK handshake complete — transport encrypted");

        Ok(Self {
            transport: Mutex::new(transport),
        })
    }

    /// Perform the Noise NK handshake as the **responder** (relay server side).
    ///
    /// `private_key` — 32-byte Curve25519 private key for the relay server.
    pub async fn handshake_responder<S>(
        private_key: &[u8; 32],
        ws: &mut S,
    ) -> Result<Self, NoiseError>
    where
        S: NoiseTransport,
    {
        let params: snow::params::NoiseParams = NOISE_PARAMS
            .parse()
            .map_err(|e: snow::Error| NoiseError::Init(e.to_string()))?;

        let mut responder = snow::Builder::new(params)
            .local_private_key(private_key)
            .build_responder()
            .map_err(|e| NoiseError::Init(e.to_string()))?;

        let mut buf = vec![0u8; 256];

        // Message 1: <- e, es
        let msg1 = ws.recv_raw().await?;
        responder
            .read_message(&msg1, &mut buf)
            .map_err(|e| NoiseError::Handshake(e.to_string()))?;

        // Message 2: -> e, ee
        let len = responder
            .write_message(&[], &mut buf)
            .map_err(|e| NoiseError::Handshake(e.to_string()))?;
        ws.send_raw(&buf[..len]).await?;

        let transport = responder
            .into_transport_mode()
            .map_err(|e| NoiseError::Handshake(e.to_string()))?;

        tracing::info!("Noise NK responder handshake complete");

        Ok(Self {
            transport: Mutex::new(transport),
        })
    }

    /// Encrypt `plaintext` and return ciphertext with AESGCM tag appended.
    ///
    /// Max plaintext size is [`NOISE_MAX_MSG`] bytes.
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseError> {
        let mut transport = self.transport.lock().unwrap();
        let mut buf = vec![0u8; plaintext.len() + 16]; // AESGCM tag is 16 bytes
        let len = transport
            .write_message(plaintext, &mut buf)
            .map_err(|e| NoiseError::Encrypt(e.to_string()))?;
        buf.truncate(len);
        Ok(buf)
    }

    /// Decrypt `ciphertext` and return plaintext.
    pub fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseError> {
        let mut transport = self.transport.lock().unwrap();
        let mut buf = vec![0u8; ciphertext.len()];
        let len = transport
            .read_message(ciphertext, &mut buf)
            .map_err(|e| NoiseError::Decrypt(e.to_string()))?;
        buf.truncate(len);
        Ok(buf)
    }
}

/// Decode a 64-character hex string into a 32-byte Curve25519 key.
fn decode_pubkey(hex: &str) -> Result<Vec<u8>, NoiseError> {
    if hex.len() != 64 {
        return Err(NoiseError::Key(format!(
            "server pubkey must be 64 hex chars (32 bytes), got {} chars",
            hex.len()
        )));
    }
    (0..hex.len())
        .step_by(2)
        .map(|i| {
            u8::from_str_radix(&hex[i..i + 2], 16)
                .map_err(|_| NoiseError::Key(format!("invalid hex at position {i}")))
        })
        .collect()
}

/// Generate a new Curve25519 keypair for the relay server.
///
/// Returns `(private_key_hex, public_key_hex)`.
pub fn generate_keypair() -> Result<(String, String), NoiseError> {
    let params: snow::params::NoiseParams = NOISE_PARAMS
        .parse()
        .map_err(|e: snow::Error| NoiseError::Init(e.to_string()))?;

    let builder = snow::Builder::new(params);
    let keypair = builder
        .generate_keypair()
        .map_err(|e| NoiseError::Init(e.to_string()))?;

    let private_hex = keypair.private.iter().map(|b| format!("{b:02x}")).collect();
    let public_hex = keypair.public.iter().map(|b| format!("{b:02x}")).collect();

    Ok((private_hex, public_hex))
}

/// Abstract transport trait for Noise handshake message exchange.
///
/// Implemented by the WebSocket write half during connection setup.
#[async_trait::async_trait]
pub trait NoiseTransport {
    async fn send_raw(&mut self, data: &[u8]) -> Result<(), NoiseError>;
    async fn recv_raw(&mut self) -> Result<Vec<u8>, NoiseError>;
}

/// Errors that can occur during Noise operations.
#[derive(Debug, thiserror::Error)]
pub enum NoiseError {
    #[error("Noise init error: {0}")]
    Init(String),
    #[error("Noise handshake error: {0}")]
    Handshake(String),
    #[error("Noise encrypt error: {0}")]
    Encrypt(String),
    #[error("Noise decrypt error: {0}")]
    Decrypt(String),
    #[error("Key error: {0}")]
    Key(String),
    #[error("Transport error: {0}")]
    Transport(String),
}

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

    #[test]
    fn test_generate_keypair() {
        let (priv_hex, pub_hex) = generate_keypair().unwrap();
        assert_eq!(priv_hex.len(), 64, "private key must be 32 bytes hex");
        assert_eq!(pub_hex.len(), 64, "public key must be 32 bytes hex");
        // Keys must be different
        assert_ne!(priv_hex, pub_hex);
    }

    #[test]
    fn test_decode_pubkey_valid() {
        let (_, pub_hex) = generate_keypair().unwrap();
        let key = decode_pubkey(&pub_hex).unwrap();
        assert_eq!(key.len(), 32);
    }

    #[test]
    fn test_decode_pubkey_invalid_length() {
        let result = decode_pubkey("deadbeef"); // too short
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("32 bytes"));
    }

    #[test]
    fn test_decode_pubkey_invalid_hex() {
        let invalid = "zz".repeat(32); // 64 chars but invalid hex
        let result = decode_pubkey(&invalid);
        assert!(result.is_err());
    }

    /// Full Noise NK round-trip: handshake + encrypt/decrypt using in-memory transport.
    #[tokio::test]
    async fn test_noise_nk_roundtrip() {
        use std::collections::VecDeque;
        use std::sync::{Arc, Mutex};

        /// In-memory transport for testing: two queues (initiator→responder, responder→initiator)
        struct MemTransport {
            send_queue: Arc<Mutex<VecDeque<Vec<u8>>>>,
            recv_queue: Arc<Mutex<VecDeque<Vec<u8>>>>,
        }

        #[async_trait::async_trait]
        impl NoiseTransport for MemTransport {
            async fn send_raw(&mut self, data: &[u8]) -> Result<(), NoiseError> {
                self.send_queue.lock().unwrap().push_back(data.to_vec());
                Ok(())
            }
            async fn recv_raw(&mut self) -> Result<Vec<u8>, NoiseError> {
                loop {
                    if let Some(msg) = self.recv_queue.lock().unwrap().pop_front() {
                        return Ok(msg);
                    }
                    tokio::task::yield_now().await;
                }
            }
        }

        // Generate server keypair
        let (priv_hex, pub_hex) = generate_keypair().unwrap();
        let priv_bytes: Vec<u8> = (0..priv_hex.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&priv_hex[i..i + 2], 16).unwrap())
            .collect();
        let priv_arr: [u8; 32] = priv_bytes.try_into().unwrap();

        // Shared queues between initiator and responder
        let i_to_r: Arc<Mutex<VecDeque<Vec<u8>>>> = Arc::new(Mutex::new(VecDeque::new()));
        let r_to_i: Arc<Mutex<VecDeque<Vec<u8>>>> = Arc::new(Mutex::new(VecDeque::new()));

        let mut initiator_transport = MemTransport {
            send_queue: i_to_r.clone(),
            recv_queue: r_to_i.clone(),
        };
        let mut responder_transport = MemTransport {
            send_queue: r_to_i.clone(),
            recv_queue: i_to_r.clone(),
        };

        // Run handshake concurrently
        let pub_hex_clone = pub_hex.clone();
        let (initiator_result, responder_result) = tokio::join!(
            NoiseSession::handshake_initiator(&pub_hex_clone, &mut initiator_transport),
            NoiseSession::handshake_responder(&priv_arr, &mut responder_transport),
        );

        let initiator = initiator_result.expect("initiator handshake");
        let responder = responder_result.expect("responder handshake");

        // Encrypt from initiator, decrypt at responder
        let plaintext = b"Hello from device to relay!";
        let ciphertext = initiator.encrypt(plaintext).unwrap();
        assert_ne!(&ciphertext, plaintext, "ciphertext must differ from plaintext");
        let recovered = responder.decrypt(&ciphertext).unwrap();
        assert_eq!(recovered, plaintext, "decrypted must match original");

        // Encrypt from responder, decrypt at initiator
        let plaintext2 = b"Response from relay to device";
        let ciphertext2 = responder.encrypt(plaintext2).unwrap();
        let recovered2 = initiator.decrypt(&ciphertext2).unwrap();
        assert_eq!(recovered2, plaintext2);
    }

    #[test]
    fn test_noise_params_are_valid() {
        // Verify the params string is accepted by snow without panicking
        let params: snow::params::NoiseParams = NOISE_PARAMS.parse().unwrap();
        assert!(!format!("{params:?}").is_empty());
    }
}
