//! Mock LLM Server — a minimal HTTP+SSE server for testing LLM channel behavior.
//!
//! Purpose: replace real LLM calls in tests with a scriptable local server that
//! speaks the OpenAI `/v1/chat/completions` protocol. Enables regression tests
//! for error classification (402/429/5xx), SSE edge cases (zero chunks, malformed),
//! and streaming semantics without any network dependency.
//!
//! Design:
//! - Pure tokio (no hyper dependency) — raw TCP + manual HTTP parse + SSE emit
//! - Each instance listens on a random localhost port
//! - Scripts are queued; each incoming request pops one script
//! - If queue is empty when a request arrives, returns 500 "no script queued"
//! - Server shuts down automatically when dropped
//!
//! Usage:
//! ```ignore
//! let server = MockLlmServer::start().await;
//! server.push_script(Script::sse_content_only("hello world")).await;
//!
//! let channel = OpenAiChannel::new(server.base_url(), "test-model", "fake-key");
//! let result = /* call stream_infer with channel */;
//! ```

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::Mutex;

/// A scripted response for a single incoming request.
#[derive(Clone)]
pub enum Script {
    /// Return HTTP 200 with an SSE stream of content chunks, then [DONE].
    SseContent { chunks: Vec<String>, finish_reason: String },
    /// Return HTTP 200 with an SSE stream containing usage info and [DONE].
    SseFull {
        content_chunks: Vec<String>,
        usage: Option<SseUsageSpec>,
        finish_reason: String,
    },
    /// Return HTTP 200 but SSE stream contains ONLY [DONE] — no content chunks.
    /// Simulates the "empty content" failure mode seen in production.
    SseDoneNoChunks,
    /// Send HTTP 200 + SSE head + content chunks, then abruptly close the
    /// connection WITHOUT sending finish_reason or [DONE]. Simulates the
    /// server-side TCP reset / network drop mid-stream failure mode.
    /// StreamInfer will see clean_eof=false → Truncated.
    SseAbruptClose { chunks: Vec<String> },
    /// Return a non-200 HTTP response with given status and body.
    /// Body is returned as raw text (typical OpenAI-compatible error JSON).
    HttpError { status: u16, body: String },
    /// Keep the connection open forever (until drop) — simulates a hang.
    /// Used to test the 60s/120s timeout paths in spawn_sse_reader.
    HangForever,
    /// Accept the request but return malformed SSE (garbage lines) before [DONE].
    MalformedSse { raw_body: String },
    /// Reflective: inspect the captured request and dynamically produce a Script.
    /// Enables happy-path tests where the response must embed the token that
    /// infer() generated into its prompt. The closure is called with the
    /// already-captured request, and returns a concrete Script (typically
    /// SseContent) built from extracted token.
    Reflective(std::sync::Arc<dyn Fn(&CapturedRequest) -> Script + Send + Sync>),
}

impl std::fmt::Debug for Script {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::SseContent { chunks, finish_reason } => f.debug_struct("SseContent")
                .field("chunks", chunks).field("finish_reason", finish_reason).finish(),
            Self::SseFull { content_chunks, usage, finish_reason } => f.debug_struct("SseFull")
                .field("content_chunks", content_chunks)
                .field("usage", usage)
                .field("finish_reason", finish_reason).finish(),
            Self::SseDoneNoChunks => write!(f, "SseDoneNoChunks"),
            Self::SseAbruptClose { chunks } => f.debug_struct("SseAbruptClose")
                .field("chunks", chunks).finish(),
            Self::HttpError { status, body } => f.debug_struct("HttpError")
                .field("status", status).field("body", body).finish(),
            Self::HangForever => write!(f, "HangForever"),
            Self::MalformedSse { raw_body } => f.debug_struct("MalformedSse")
                .field("raw_body", raw_body).finish(),
            Self::Reflective(_) => write!(f, "Reflective(<closure>)"),
        }
    }
}

#[derive(Debug, Clone)]
pub struct SseUsageSpec {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub cached_tokens: u64,
}

impl Script {
    /// Convenience: simple success with only content chunks and finish_reason=stop.
    pub fn sse_content_only(text: &str) -> Self {
        Script::SseContent {
            chunks: vec![text.to_string()],
            finish_reason: "stop".to_string(),
        }
    }

    /// Convenience: split text into N char-sized chunks for streaming simulation.
    pub fn sse_streamed(text: &str, chunks: usize) -> Self {
        let chars: Vec<char> = text.chars().collect();
        let chunk_size = (chars.len() + chunks - 1) / chunks.max(1);
        let chunks: Vec<String> = chars
            .chunks(chunk_size.max(1))
            .map(|c| c.iter().collect())
            .collect();
        Script::SseContent {
            chunks,
            finish_reason: "stop".to_string(),
        }
    }

    /// Convenience: HTTP 402 quota-exceeded (OpenAI-style error body).
    pub fn quota_exceeded() -> Self {
        Script::HttpError {
            status: 402,
            body: r#"{"error":{"message":"You exceeded your current quota","type":"insufficient_quota","code":"insufficient_quota"}}"#.to_string(),
        }
    }

    /// Convenience: HTTP 429 rate limited.
    pub fn rate_limited() -> Self {
        Script::HttpError {
            status: 429,
            body: r#"{"error":{"message":"Rate limit reached","type":"requests","code":"rate_limit_exceeded"}}"#.to_string(),
        }
    }

    /// Convenience: HTTP 500 server error.
    pub fn server_error() -> Self {
        Script::HttpError {
            status: 500,
            body: r#"{"error":{"message":"Internal server error"}}"#.to_string(),
        }
    }
}

/// The mock server handle.
pub struct MockLlmServer {
    addr: SocketAddr,
    scripts: Arc<Mutex<Vec<Script>>>,
    requests: Arc<Mutex<Vec<CapturedRequest>>>,
    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
}

/// A captured incoming request (for assertions).
#[derive(Debug, Clone)]
pub struct CapturedRequest {
    pub method: String,
    pub path: String,
    pub authorization: Option<String>,
    pub body: String,
}

impl MockLlmServer {
    /// Start a new mock server on a random localhost port.
    pub async fn start() -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
        let addr = listener.local_addr().expect("local_addr");
        let scripts = Arc::new(Mutex::new(Vec::<Script>::new()));
        let requests = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));

        let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();

        let scripts_c = scripts.clone();
        let requests_c = requests.clone();

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = &mut shutdown_rx => break,
                    accept = listener.accept() => {
                        let (socket, _) = match accept {
                            Ok(x) => x,
                            Err(_) => continue,
                        };
                        let scripts = scripts_c.clone();
                        let requests = requests_c.clone();
                        tokio::spawn(async move {
                            handle_connection(socket, scripts, requests).await;
                        });
                    }
                }
            }
        });

        Self {
            addr,
            scripts,
            requests,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// Base URL of the server, e.g. `http://127.0.0.1:PORT`.
    pub fn base_url(&self) -> String {
        format!("http://{}", self.addr)
    }

    /// Full chat completions URL, e.g. `http://127.0.0.1:PORT/v1/chat/completions`.
    pub fn chat_url(&self) -> String {
        format!("http://{}/v1/chat/completions", self.addr)
    }

    /// Queue a script for the next incoming request.
    pub async fn push_script(&self, script: Script) {
        self.scripts.lock().await.push(script);
    }

    /// Queue multiple scripts (consumed in FIFO order).
    pub async fn push_scripts(&self, scripts: Vec<Script>) {
        self.scripts.lock().await.extend(scripts);
    }

    /// Retrieve all captured requests.
    pub async fn requests(&self) -> Vec<CapturedRequest> {
        self.requests.lock().await.clone()
    }

    /// Number of requests received so far.
    pub async fn request_count(&self) -> usize {
        self.requests.lock().await.len()
    }
}

impl Drop for MockLlmServer {
    fn drop(&mut self) {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
    }
}

// ============================================================
// Connection handler — minimal HTTP/1.1 parser + SSE emitter
// ============================================================

async fn handle_connection(
    mut socket: tokio::net::TcpStream,
    scripts: Arc<Mutex<Vec<Script>>>,
    requests: Arc<Mutex<Vec<CapturedRequest>>>,
) {
    // 1. Read request: headers until \r\n\r\n, then body based on Content-Length.
    let mut buf = Vec::<u8>::with_capacity(4096);
    let mut chunk = [0u8; 1024];
    let header_end = loop {
        let n = match socket.read(&mut chunk).await {
            Ok(0) => return, // EOF before request complete
            Ok(n) => n,
            Err(_) => return,
        };
        buf.extend_from_slice(&chunk[..n]);
        if let Some(pos) = find_subslice(&buf, b"\r\n\r\n") {
            break pos + 4;
        }
        if buf.len() > 1_000_000 {
            return; // protect against runaway request
        }
    };

    let header_text = String::from_utf8_lossy(&buf[..header_end - 4]).to_string();
    let (method, path, headers) = parse_request_head(&header_text);

    // Determine body length
    let content_length: usize = headers
        .iter()
        .find(|(k, _)| k.eq_ignore_ascii_case("content-length"))
        .and_then(|(_, v)| v.trim().parse::<usize>().ok())
        .unwrap_or(0);

    let mut body = buf[header_end..].to_vec();
    while body.len() < content_length {
        let n = match socket.read(&mut chunk).await {
            Ok(0) => break,
            Ok(n) => n,
            Err(_) => return,
        };
        body.extend_from_slice(&chunk[..n]);
    }
    body.truncate(content_length);

    let authorization = headers
        .iter()
        .find(|(k, _)| k.eq_ignore_ascii_case("authorization"))
        .map(|(_, v)| v.trim().to_string());

    let body_str = String::from_utf8_lossy(&body).to_string();

    requests.lock().await.push(CapturedRequest {
        method: method.clone(),
        path: path.clone(),
        authorization,
        body: body_str,
    });

    // 2. Pop next script and respond.
    let script = {
        let mut q = scripts.lock().await;
        if q.is_empty() {
            None
        } else {
            Some(q.remove(0))
        }
    };

    let script = match script {
        Some(s) => s,
        None => {
            let body = b"no script queued in MockLlmServer";
            let resp = format!(
                "HTTP/1.1 500 Internal Server Error\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                body.len()
            );
            let _ = socket.write_all(resp.as_bytes()).await;
            let _ = socket.write_all(body).await;
            return;
        }
    };

    // Resolve Reflective scripts: call the closure with the just-captured
    // request, then dispatch the returned concrete script. We re-fetch the
    // last captured request (the one that triggered this dispatch).
    let resolved = match script {
        Script::Reflective(f) => {
            let captured = {
                let q = requests.lock().await;
                q.last().cloned().expect("captured request must exist before dispatch")
            };
            f(&captured)
        }
        other => other,
    };

    execute_script(&mut socket, resolved).await;
}

async fn execute_script(socket: &mut tokio::net::TcpStream, script: Script) {
    match script {
        Script::SseContent { chunks, finish_reason } => {
            send_sse_head(socket).await;
            for c in &chunks {
                send_sse_data_line(socket, &sse_content_json(c)).await;
            }
            send_sse_data_line(socket, &sse_finish_json(&finish_reason)).await;
            send_sse_done(socket).await;
        }
        Script::SseFull { content_chunks, usage, finish_reason } => {
            send_sse_head(socket).await;
            for c in &content_chunks {
                send_sse_data_line(socket, &sse_content_json(c)).await;
            }
            send_sse_data_line(socket, &sse_finish_json(&finish_reason)).await;
            if let Some(u) = usage {
                send_sse_data_line(socket, &sse_usage_json(&u)).await;
            }
            send_sse_done(socket).await;
        }
        Script::SseDoneNoChunks => {
            send_sse_head(socket).await;
            send_sse_done(socket).await;
        }
        Script::SseAbruptClose { chunks } => {
            send_sse_head(socket).await;
            for c in &chunks {
                send_sse_data_line(socket, &sse_content_json(c)).await;
            }
            // Intentionally skip finish_reason + [DONE]; shutdown below closes socket.
        }
        Script::HttpError { status, body } => {
            let reason = status_reason(status);
            let resp = format!(
                "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                status,
                reason,
                body.len()
            );
            let _ = socket.write_all(resp.as_bytes()).await;
            let _ = socket.write_all(body.as_bytes()).await;
        }
        Script::HangForever => {
            // Hold the connection open. When the test drops the server, the
            // tokio task ends and the socket is closed.
            send_sse_head(socket).await;
            tokio::time::sleep(Duration::from_secs(3600)).await;
        }
        Script::MalformedSse { raw_body } => {
            send_sse_head(socket).await;
            let _ = socket.write_all(raw_body.as_bytes()).await;
            send_sse_done(socket).await;
        }
        Script::Reflective(_) => {
            unreachable!("Reflective must be resolved by dispatcher before execute_script")
        }
    }
    let _ = socket.shutdown().await;
}

async fn send_sse_head(socket: &mut tokio::net::TcpStream) {
    let head = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n";
    let _ = socket.write_all(head.as_bytes()).await;
}

async fn send_sse_data_line(socket: &mut tokio::net::TcpStream, json_body: &str) {
    let line = format!("data: {}\n\n", json_body);
    let _ = socket.write_all(line.as_bytes()).await;
}

async fn send_sse_done(socket: &mut tokio::net::TcpStream) {
    let line = "data: [DONE]\n\n";
    let _ = socket.write_all(line.as_bytes()).await;
}

fn sse_content_json(content: &str) -> String {
    let escaped = json_escape(content);
    format!(
        r#"{{"choices":[{{"delta":{{"content":"{}"}},"finish_reason":null}}]}}"#,
        escaped
    )
}

fn sse_finish_json(reason: &str) -> String {
    format!(
        r#"{{"choices":[{{"delta":{{}},"finish_reason":"{}"}}]}}"#,
        reason
    )
}

fn sse_usage_json(u: &SseUsageSpec) -> String {
    format!(
        r#"{{"choices":[],"usage":{{"prompt_tokens":{},"completion_tokens":{},"prompt_tokens_details":{{"cached_tokens":{}}}}}}}"#,
        u.prompt_tokens, u.completion_tokens, u.cached_tokens
    )
}

fn json_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
            c => out.push(c),
        }
    }
    out
}

fn status_reason(status: u16) -> &'static str {
    match status {
        200 => "OK",
        400 => "Bad Request",
        401 => "Unauthorized",
        402 => "Payment Required",
        403 => "Forbidden",
        404 => "Not Found",
        408 => "Request Timeout",
        429 => "Too Many Requests",
        500 => "Internal Server Error",
        502 => "Bad Gateway",
        503 => "Service Unavailable",
        504 => "Gateway Timeout",
        _ => "Error",
    }
}

fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    if needle.is_empty() || needle.len() > haystack.len() {
        return None;
    }
    (0..=haystack.len() - needle.len()).find(|&i| &haystack[i..i + needle.len()] == needle)
}

fn parse_request_head(text: &str) -> (String, String, Vec<(String, String)>) {
    let mut lines = text.split("\r\n");
    let first = lines.next().unwrap_or("");
    let mut parts = first.split_whitespace();
    let method = parts.next().unwrap_or("").to_string();
    let path = parts.next().unwrap_or("").to_string();
    let headers: Vec<(String, String)> = lines
        .filter_map(|l| {
            let idx = l.find(':')?;
            Some((l[..idx].to_string(), l[idx + 1..].to_string()))
        })
        .collect();
    (method, path, headers)
}