selfware 0.6.3

//! Mock LLM API Server for CI Testing
//!
//! Provides a [`MockLlmServer`] that emulates an OpenAI-compatible
//! `/v1/chat/completions` endpoint. Designed for deterministic unit and
//! integration tests that must not depend on a live model endpoint.
//!
//! # Features
//! - Canned text responses
//! - Tool-call responses (function calling)
//! - Configurable error responses (status code + body)
//! - Latency simulation
//! - Builder pattern for ergonomic test setup
//!
//! # Example
//! ```ignore
//! let server = MockLlmServer::builder()
//!     .with_response("Hello from mock!")
//!     .with_latency_ms(50)
//!     .build()
//!     .await;
//! let url = server.url(); // e.g. "http://127.0.0.1:12345"
//! // ... point your client at `url` ...
//! server.stop().await;
//! ```

use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::{watch, Mutex};

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// A single tool call that the mock server can include in its response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MockToolCall {
    /// Unique call id, e.g. `"call_abc123"`.
    pub id: String,
    /// Tool function name, e.g. `"file_read"`.
    pub name: String,
    /// JSON-encoded arguments string, e.g. `"{\"path\":\"foo.rs\"}"`.
    pub arguments: String,
}

/// Describes how the mock server should respond to the next request.
#[derive(Debug, Clone)]
pub enum MockResponse {
    /// Return a plain assistant text message.
    Text(String),
    /// Return a response containing tool calls.
    ToolCalls(Vec<MockToolCall>),
    /// Return an HTTP error with the given status code and body.
    Error { status: u16, body: String },
    /// Return an HTTP error with custom headers (e.g. Retry-After).
    ErrorWithHeaders {
        status: u16,
        body: String,
        headers: Vec<(String, String)>,
    },
}

/// A lightweight mock HTTP server that speaks just enough of the
/// OpenAI chat-completions protocol to satisfy selfware's API client.
pub struct MockLlmServer {
    /// Local URL the server listens on, e.g. `"http://127.0.0.1:PORT"`.
    url: String,
    /// Sender half of a shutdown signal.
    shutdown_tx: watch::Sender<bool>,
    /// Join handle for the background accept loop.
    handle: tokio::task::JoinHandle<()>,
    /// Captured request bodies for assertion in tests.
    captured_requests: Arc<Mutex<Vec<String>>>,
}

impl MockLlmServer {
    /// Create a new [`MockLlmServerBuilder`] for ergonomic configuration.
    pub fn builder() -> MockLlmServerBuilder {
        MockLlmServerBuilder::default()
    }

    /// Start the mock server with the given configuration.
    ///
    /// Binds to `127.0.0.1:0` (OS-assigned port), spawns a background
    /// tokio task, and returns a handle that can be used to query the URL
    /// or stop the server.
    pub async fn start(config: MockServerConfig) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("failed to bind mock server");
        let addr = listener.local_addr().expect("failed to get local addr");
        let url = format!("http://{}", addr);

        let (shutdown_tx, shutdown_rx) = watch::channel(false);
        let config = Arc::new(config);
        let captured_requests = Arc::new(Mutex::new(Vec::new()));

        let handle = tokio::spawn(accept_loop(
            listener,
            config,
            shutdown_rx,
            Arc::clone(&captured_requests),
        ));

        Self {
            url,
            shutdown_tx,
            handle,
            captured_requests,
        }
    }

    /// The base URL of this mock server (e.g. `"http://127.0.0.1:54321"`).
    pub fn url(&self) -> &str {
        &self.url
    }

    /// Return the JSON request bodies captured so far (in order received).
    pub async fn captured_request_bodies(&self) -> Vec<String> {
        self.captured_requests.lock().await.clone()
    }

    /// Signal the server to stop accepting new connections and wait for the
    /// background task to finish.
    pub async fn stop(self) {
        let _ = self.shutdown_tx.send(true);
        let _ = self.handle.await;
    }
}

// ---------------------------------------------------------------------------
// Builder
// ---------------------------------------------------------------------------

/// Configuration produced by the builder that drives mock server behaviour.
#[derive(Debug, Clone)]
pub struct MockServerConfig {
    /// Queue of responses. Each request pops the next response; when the
    /// queue is exhausted the server falls back to `default_response`.
    pub responses: Vec<MockResponse>,
    /// Response used after the queue is empty.
    pub default_response: MockResponse,
    /// Artificial latency added before every response (milliseconds).
    pub latency_ms: u64,
    /// Model name included in the JSON response body.
    pub model: String,
}

impl Default for MockServerConfig {
    fn default() -> Self {
        Self {
            responses: Vec::new(),
            default_response: MockResponse::Text("Hello from MockLlmServer".to_string()),
            latency_ms: 0,
            model: "mock-model".to_string(),
        }
    }
}

/// Builder for [`MockLlmServer`] providing a fluent configuration API.
#[derive(Default)]
pub struct MockLlmServerBuilder {
    config: MockServerConfig,
}

impl MockLlmServerBuilder {
    /// Queue a plain text response. Responses are served in FIFO order.
    pub fn with_response(mut self, text: impl Into<String>) -> Self {
        self.config.responses.push(MockResponse::Text(text.into()));
        self
    }

    /// Queue a tool-call response.
    pub fn with_tool_calls(mut self, calls: Vec<MockToolCall>) -> Self {
        self.config.responses.push(MockResponse::ToolCalls(calls));
        self
    }

    /// Queue an error response with the specified HTTP status code and body.
    pub fn with_error(mut self, status: u16, body: impl Into<String>) -> Self {
        self.config.responses.push(MockResponse::Error {
            status,
            body: body.into(),
        });
        self
    }

    /// Queue an error with custom headers (e.g. 429 with Retry-After).
    pub fn with_error_and_headers(
        mut self,
        status: u16,
        body: impl Into<String>,
        headers: Vec<(String, String)>,
    ) -> Self {
        self.config.responses.push(MockResponse::ErrorWithHeaders {
            status,
            body: body.into(),
            headers,
        });
        self
    }

    /// Set the artificial latency (in milliseconds) applied before every
    /// response is sent.
    pub fn with_latency(mut self, ms: u64) -> Self {
        self.config.latency_ms = ms;
        self
    }

    /// Override the model name returned in response bodies.
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.config.model = model.into();
        self
    }

    /// Set the fallback response used once the queued responses are exhausted.
    pub fn with_default_response(mut self, resp: MockResponse) -> Self {
        self.config.default_response = resp;
        self
    }

    /// Build and start the mock server, returning a ready-to-use handle.
    pub async fn build(self) -> MockLlmServer {
        MockLlmServer::start(self.config).await
    }
}

// ---------------------------------------------------------------------------
// Internal: accept loop & request handling
// ---------------------------------------------------------------------------

/// Background accept loop. Runs until `shutdown_rx` signals true.
async fn accept_loop(
    listener: TcpListener,
    config: Arc<MockServerConfig>,
    mut shutdown_rx: watch::Receiver<bool>,
    captured_requests: Arc<Mutex<Vec<String>>>,
) {
    let response_idx = Arc::new(Mutex::new(0usize));

    loop {
        tokio::select! {
            _ = shutdown_rx.changed() => {
                if *shutdown_rx.borrow() {
                    break;
                }
            }
            accept_result = listener.accept() => {
                match accept_result {
                    Ok((stream, _addr)) => {
                        let cfg = Arc::clone(&config);
                        let idx = Arc::clone(&response_idx);
                        let cap = Arc::clone(&captured_requests);
                        tokio::spawn(async move {
                            if let Err(e) = handle_connection(stream, cfg, idx, cap).await {
                                tracing::debug!("mock server connection error: {}", e);
                            }
                        });
                    }
                    Err(e) => {
                        tracing::debug!("mock server accept error: {}", e);
                    }
                }
            }
        }
    }
}

/// Handle a single HTTP connection. Reads the request, determines the
/// response from the config, and writes raw HTTP back.
async fn handle_connection(
    mut stream: tokio::net::TcpStream,
    config: Arc<MockServerConfig>,
    response_idx: Arc<Mutex<usize>>,
    captured_requests: Arc<Mutex<Vec<String>>>,
) -> std::io::Result<()> {
    let mut buf = vec![0u8; 65536];
    let n = stream.read(&mut buf).await?;
    if n == 0 {
        return Ok(());
    }

    let request = String::from_utf8_lossy(&buf[..n]);

    // Capture the JSON body (everything after the blank line separator).
    if let Some(body_start) = request.find("\r\n\r\n") {
        let body = request[body_start + 4..].trim().to_string();
        if !body.is_empty() {
            captured_requests.lock().await.push(body);
        }
    }

    // Handle POST /v1/chat/completions and /v1/completions
    let is_post = request.starts_with("POST");
    let is_chat = is_post && request.contains("/v1/chat/completions");
    let is_completion = is_post && request.contains("/v1/completions") && !is_chat;
    // A streaming request must be answered with SSE, not a plain JSON body —
    // otherwise the client's streaming parser gets no events and the run stalls.
    let is_streaming = request.contains("\"stream\":true") || request.contains("\"stream\": true");

    if !is_chat && !is_completion {
        let response = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n";
        stream.write_all(response.as_bytes()).await?;
        return Ok(());
    }

    // Apply configured latency
    if config.latency_ms > 0 {
        tokio::time::sleep(std::time::Duration::from_millis(config.latency_ms)).await;
    }

    // Pick the next response
    let mock_response = {
        let mut idx = response_idx.lock().await;
        if *idx < config.responses.len() {
            let resp = config.responses[*idx].clone();
            *idx += 1;
            resp
        } else {
            config.default_response.clone()
        }
    };

    match mock_response {
        MockResponse::Text(text) => {
            if is_streaming {
                write_sse_text_response(&mut stream, &text).await?;
            } else {
                let body = format_chat_response(&config.model, &text, None);
                write_http_response(&mut stream, 200, &body, &[]).await?;
            }
        }
        MockResponse::ToolCalls(calls) => {
            let tool_calls_json = format_tool_calls(&calls);
            let body = format_chat_response(&config.model, "", Some(&tool_calls_json));
            write_http_response(&mut stream, 200, &body, &[]).await?;
        }
        MockResponse::Error { status, body } => {
            write_http_response(&mut stream, status, &body, &[]).await?;
        }
        MockResponse::ErrorWithHeaders {
            status,
            body,
            headers,
        } => {
            let header_refs: Vec<(&str, &str)> = headers
                .iter()
                .map(|(k, v)| (k.as_str(), v.as_str()))
                .collect();
            write_http_response(&mut stream, status, &body, &header_refs).await?;
        }
    }

    Ok(())
}

/// Format a JSON body conforming to the OpenAI chat completions response
/// schema.
fn format_chat_response(model: &str, content: &str, tool_calls: Option<&str>) -> String {
    let tool_calls_field = match tool_calls {
        Some(tc) => format!(r#","tool_calls":{}"#, tc),
        None => String::new(),
    };

    let finish_reason = if tool_calls.is_some() {
        "tool_calls"
    } else {
        "stop"
    };

    // Escape content for JSON embedding
    let escaped_content = serde_json::to_string(content).unwrap_or_else(|_| "\"\"".to_string());
    // Remove surrounding quotes since we embed it in the template
    let escaped_content = &escaped_content[1..escaped_content.len() - 1];

    format!(
        r#"{{"id":"mock-resp-1","object":"chat.completion","created":1700000000,"model":"{}","choices":[{{"index":0,"message":{{"role":"assistant","content":"{}"{}}},"finish_reason":"{}"}}],"usage":{{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}}}"#,
        model, escaped_content, tool_calls_field, finish_reason,
    )
}

/// Write a Server-Sent Events (SSE) streaming chat response: the content as a
/// single delta chunk, a final `finish_reason:"stop"` chunk with usage, then
/// `[DONE]` — chunked-transfer-encoded like a real OpenAI-compatible stream.
async fn write_sse_text_response(
    stream: &mut tokio::net::TcpStream,
    content: &str,
) -> std::io::Result<()> {
    use tokio::io::AsyncWriteExt;

    let escaped = serde_json::to_string(content).unwrap_or_else(|_| "\"\"".to_string());
    let escaped = &escaped[1..escaped.len() - 1];

    let events = format!(
        "data: {{\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"{}\"}},\"finish_reason\":null}}]}}\n\n\
         data: {{\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}],\"usage\":{{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}}}\n\n\
         data: [DONE]\n\n",
        escaped
    );
    let chunk = format!("{:X}\r\n{}\r\n", events.len(), events);
    let response = format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n{}0\r\n\r\n",
        chunk
    );
    stream.write_all(response.as_bytes()).await?;
    stream.shutdown().await
}

/// Serialize a list of [`MockToolCall`]s into the JSON array format expected
/// by the OpenAI API.
fn format_tool_calls(calls: &[MockToolCall]) -> String {
    let items: Vec<String> = calls
        .iter()
        .map(|c| {
            // Escape arguments string for safe JSON embedding
            let escaped_args =
                serde_json::to_string(&c.arguments).unwrap_or_else(|_| "\"{}\"".to_string());
            format!(
                r#"{{"id":"{}","type":"function","function":{{"name":"{}","arguments":{}}}}}"#,
                c.id, c.name, escaped_args,
            )
        })
        .collect();
    format!("[{}]", items.join(","))
}

/// Write a full HTTP/1.1 response to the stream.
async fn write_http_response(
    stream: &mut tokio::net::TcpStream,
    status: u16,
    body: &str,
    extra_headers: &[(&str, &str)],
) -> std::io::Result<()> {
    let status_text = match status {
        200 => "OK",
        400 => "Bad Request",
        401 => "Unauthorized",
        404 => "Not Found",
        429 => "Too Many Requests",
        500 => "Internal Server Error",
        503 => "Service Unavailable",
        _ => "Error",
    };

    let mut extra = String::new();
    for (key, value) in extra_headers {
        extra.push_str(&format!("{}: {}\r\n", key, value));
    }

    let response = format!(
        "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n{}Connection: close\r\n\r\n{}",
        status,
        status_text,
        body.len(),
        extra,
        body,
    );

    stream.write_all(response.as_bytes()).await
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[path = "../../tests/unit/testing/mock_api/mock_api_test.rs"]
mod tests;