// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 LLMHosts.com Contributors
//
//! LLMHosts Relay Client — Rust core compiled via pyo3 + maturin.
//!
//! M1: Basic WebSocket connection with Bearer token auth.
//! M2: + Noise NK encryption (snow crate) — zero-knowledge relay.
//! M3: + yamux stream multiplexing.
//!
//! Python usage (M2):
//!     from _relay_core import RelayClient, RELAY_VERSION
//!     client = RelayClient(
//!         "wss://llmhosts.com/api/relay/connect",
//!         "device_token",
//!         11400,
//!         server_pubkey="<64-char hex Curve25519 pubkey>",
//!     )
//!     await client.connect()  # performs Noise NK handshake
//!     assert client.is_connected
//!     await client.disconnect()
//!
//! If server_pubkey is None, falls back to M1 (no encryption). This allows
//! backward compatibility during relay server rollout.

pub mod noise;

use noise::NoiseSession;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;

/// Shared state between the Python-facing struct and the background tokio task.
struct ClientState {
    connected: AtomicBool,
    bytes_relayed: AtomicU64,
    streams_active: AtomicU64,
    connect_time: std::sync::Mutex<Option<Instant>>,
    /// Noise NK session — Some after successful handshake, None in M1 fallback.
    noise: std::sync::Mutex<Option<NoiseSession>>,
}

impl ClientState {
    fn new() -> Self {
        Self {
            connected: AtomicBool::new(false),
            bytes_relayed: AtomicU64::new(0),
            streams_active: AtomicU64::new(0),
            connect_time: std::sync::Mutex::new(None),
            noise: std::sync::Mutex::new(None),
        }
    }
}

/// Relay client that connects to the LLMHosts relay server via WebSocket.
///
/// M2 behaviour: if `server_pubkey` is provided, performs a Noise NK handshake
/// immediately after the WebSocket upgrade. All subsequent frames are encrypted
/// with AESGCM-256 — the relay server (VPS) cannot read payload contents.
///
/// If `server_pubkey` is None, operates in M1 mode (unencrypted, for
/// development and backward compatibility during relay server rollout).
#[pyclass]
struct RelayClient {
    relay_url: String,
    device_token: String,
    local_port: u16,
    /// Optional 64-char hex Curve25519 server public key for Noise NK (M2).
    server_pubkey: Option<String>,
    state: Arc<ClientState>,
}

#[pymethods]
impl RelayClient {
    /// Create a new RelayClient.
    ///
    /// Args:
    ///     relay_url: WebSocket URL of the relay server (e.g. "wss://llmhosts.com/api/relay/connect")
    ///     device_token: Authentication token for this device
    ///     local_port: Local port that the FastAPI proxy listens on
    ///     server_pubkey: Optional 64-char hex Curve25519 server public key.
    ///                    If provided, performs Noise NK handshake after WS connect (M2).
    ///                    If None, uses unencrypted M1 mode.
    #[new]
    #[pyo3(signature = (relay_url, device_token, local_port, server_pubkey=None))]
    fn new(relay_url: String, device_token: String, local_port: u16, server_pubkey: Option<String>) -> Self {
        Self {
            relay_url,
            device_token,
            local_port,
            server_pubkey,
            state: Arc::new(ClientState::new()),
        }
    }

    /// Connect to the relay server via WebSocket.
    ///
    /// If `server_pubkey` was provided at construction, performs Noise NK
    /// handshake immediately after the WebSocket upgrade — all subsequent
    /// frames will be encrypted (M2 mode).
    ///
    /// Returns when the connection (and optional handshake) is established.
    fn connect<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let url = self.relay_url.clone();
        let token = self.device_token.clone();
        let state = self.state.clone();
        let server_pubkey = self.server_pubkey.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            use futures_util::SinkExt;
            use tokio_tungstenite::connect_async;
            use tokio_tungstenite::tungstenite::http::Request;
            use tokio_tungstenite::tungstenite::Message;

            // Build WebSocket request with Authorization header
            let request = Request::builder()
                .uri(&url)
                .header("Authorization", format!("Bearer {}", token))
                .header("Host", extract_host(&url).unwrap_or_default())
                .header("Connection", "Upgrade")
                .header("Upgrade", "websocket")
                .header("Sec-WebSocket-Version", "13")
                .header(
                    "Sec-WebSocket-Key",
                    tokio_tungstenite::tungstenite::handshake::client::generate_key(),
                )
                .body(())
                .map_err(|e| PyRuntimeError::new_err(format!("Invalid request: {e}")))?;

            let (ws_stream, _response) = connect_async(request)
                .await
                .map_err(|e| PyRuntimeError::new_err(format!("WebSocket connect failed: {e}")))?;

            // M2: Noise NK handshake if server_pubkey is provided
            if let Some(ref pubkey_hex) = server_pubkey {
                use futures_util::StreamExt;
                let (mut write, mut read) = ws_stream.split();

                let params: snow::params::NoiseParams = "Noise_NK_25519_AESGCM_SHA256"
                    .parse()
                    .map_err(|e: snow::Error| PyRuntimeError::new_err(format!("Noise params: {e}")))?;
                let server_pubkey_bytes = decode_hex(pubkey_hex)
                    .map_err(|e| PyRuntimeError::new_err(format!("Invalid server_pubkey: {e}")))?;

                let mut initiator = snow::Builder::new(params)
                    .remote_public_key(&server_pubkey_bytes)
                    .build_initiator()
                    .map_err(|e| PyRuntimeError::new_err(format!("Noise init: {e}")))?;

                // -> e, es
                let mut buf = vec![0u8; 256];
                let len = initiator
                    .write_message(&[], &mut buf)
                    .map_err(|e| PyRuntimeError::new_err(format!("Noise write_message: {e}")))?;
                write
                    .send(Message::Binary(bytes::Bytes::copy_from_slice(&buf[..len])))
                    .await
                    .map_err(|e| PyRuntimeError::new_err(format!("WS send: {e}")))?;

                // <- e, ee
                let msg2 = read
                    .next()
                    .await
                    .ok_or_else(|| PyRuntimeError::new_err("Noise handshake: server closed"))?
                    .map_err(|e| PyRuntimeError::new_err(format!("Noise recv: {e}")))?;
                let msg2_data = msg2.into_data();
                initiator
                    .read_message(&msg2_data, &mut buf)
                    .map_err(|e| PyRuntimeError::new_err(format!("Noise read_message: {e}")))?;

                let transport = initiator
                    .into_transport_mode()
                    .map_err(|e| PyRuntimeError::new_err(format!("Noise transport mode: {e}")))?;

                let noise_session = NoiseSession::from_transport(transport);
                tracing::info!(relay_url = %url, "Noise NK handshake complete — relay encrypted");
                *state.noise.lock().unwrap() = Some(noise_session);

                state.connected.store(true, Ordering::SeqCst);
                *state.connect_time.lock().unwrap() = Some(Instant::now());

                // Spawn background task to keep WS alive (M3 will add yamux)
                let state_bg = state.clone();
                tokio::spawn(async move {
                    while let Some(msg) = read.next().await {
                        match msg {
                            Ok(m) => {
                                state_bg.bytes_relayed.fetch_add(m.len() as u64, Ordering::Relaxed);
                            }
                            Err(_) => {
                                state_bg.connected.store(false, Ordering::SeqCst);
                                break;
                            }
                        }
                    }
                    state_bg.connected.store(false, Ordering::SeqCst);
                });
            } else {
                // M1 fallback — no encryption
                tracing::warn!("No server_pubkey provided — relay running without Noise NK encryption (M1 mode)");
                state.connected.store(true, Ordering::SeqCst);
                *state.connect_time.lock().unwrap() = Some(Instant::now());

                let state_bg = state.clone();
                tokio::spawn(async move {
                    use futures_util::StreamExt;
                    let (_write, mut read) = ws_stream.split();
                    while let Some(msg) = read.next().await {
                        match msg {
                            Ok(m) => {
                                state_bg.bytes_relayed.fetch_add(m.len() as u64, Ordering::Relaxed);
                            }
                            Err(_) => {
                                state_bg.connected.store(false, Ordering::SeqCst);
                                break;
                            }
                        }
                    }
                    state_bg.connected.store(false, Ordering::SeqCst);
                });
            }

            Ok(())
        })
    }

    /// Disconnect from the relay server.
    fn disconnect<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let state = self.state.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            state.connected.store(false, Ordering::SeqCst);
            // M3+ will properly close the yamux session and WebSocket
            Ok(())
        })
    }

    /// Whether the client is currently connected.
    #[getter]
    fn is_connected(&self) -> bool {
        self.state.connected.load(Ordering::SeqCst)
    }

    /// The relay URL this client connects to.
    #[getter]
    fn relay_url(&self) -> &str {
        &self.relay_url
    }

    /// The local port being forwarded.
    #[getter]
    fn local_port(&self) -> u16 {
        self.local_port
    }

    /// Total bytes relayed through this connection.
    #[getter]
    fn bytes_relayed(&self) -> u64 {
        self.state.bytes_relayed.load(Ordering::Relaxed)
    }

    /// Number of active yamux streams (M3+, returns 0 for M1).
    #[getter]
    fn streams_active(&self) -> u64 {
        self.state.streams_active.load(Ordering::Relaxed)
    }

    /// Connection uptime in seconds.
    #[getter]
    fn uptime_seconds(&self) -> f64 {
        let guard = self.state.connect_time.lock().unwrap();
        match *guard {
            Some(t) if self.state.connected.load(Ordering::SeqCst) => t.elapsed().as_secs_f64(),
            _ => 0.0,
        }
    }

    /// Whether Noise NK encryption is active on this connection (M2 mode).
    ///
    /// Returns True only after a successful Noise handshake — confirming that
    /// all relay traffic is end-to-end encrypted and the VPS cannot read payloads.
    #[getter]
    fn is_noise_encrypted(&self) -> bool {
        self.state.noise.lock().unwrap().is_some()
    }
}

/// Extract host from a URL string.
fn extract_host(url: &str) -> Option<String> {
    url::Url::parse(url).ok().and_then(|u| u.host_str().map(|h| h.to_string()))
}

/// Decode a hex string into bytes.
fn decode_hex(hex: &str) -> Result<Vec<u8>, String> {
    if !hex.len().is_multiple_of(2) {
        return Err(format!("odd hex length: {}", hex.len()));
    }
    (0..hex.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).map_err(|_| format!("invalid hex at {i}")))
        .collect()
}

pub const RELAY_VERSION: &str = env!("CARGO_PKG_VERSION");

/// Generate a Curve25519 keypair for the relay server.
///
/// Returns `(private_key_hex, public_key_hex)`. Run this on the VPS once,
/// store private_key_hex in the server environment, and publish public_key_hex
/// so devices can verify the server's identity.
#[pyfunction]
fn generate_relay_keypair() -> PyResult<(String, String)> {
    noise::generate_keypair()
        .map_err(|e| PyRuntimeError::new_err(format!("Keypair generation failed: {e}")))
}

/// Register relay_core classes/constants into a parent module.
pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<RelayClient>()?;
    m.add("RELAY_VERSION", RELAY_VERSION)?;
    m.add("RUST_AVAILABLE", true)?;
    m.add_function(wrap_pyfunction!(generate_relay_keypair, m)?)?;
    Ok(())
}

/// Python module definition.
#[pymodule]
fn _relay_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    register_module(m)
}

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

    #[test]
    fn test_extract_host() {
        assert_eq!(
            extract_host("wss://llmhosts.com/api/relay/connect"),
            Some("llmhosts.com".to_string())
        );
        assert_eq!(
            extract_host("ws://localhost:8080/ws"),
            Some("localhost".to_string())
        );
        assert_eq!(extract_host("not-a-url"), None);
    }
}
