rust-sanitize 0.14.1

use std::io::{BufRead, BufReader, Read, Write};
use std::time::Duration;

/// Maximum bytes accepted from the SSE response stream.
const MAX_STREAM_BYTES: usize = 10 * 1024 * 1024; // 10 MB

/// Maximum bytes read from an HTTP error body.
const MAX_ERROR_BODY_BYTES: u64 = 4 * 1024; // 4 KB

/// Validate that `endpoint` uses an http or https scheme.
/// Called from `validate_args` before the HTTP request is made.
pub(crate) fn validate_endpoint_scheme(endpoint: &str) -> Result<(), String> {
    if !endpoint.starts_with("http://") && !endpoint.starts_with("https://") {
        return Err(format!(
            "--llm-endpoint must start with http:// or https://; got: {endpoint}"
        ));
    }
    Ok(())
}

/// Remove ESC bytes from `s` to prevent terminal control-sequence injection.
fn strip_terminal_escapes(s: &str) -> String {
    s.replace('\x1b', "")
}

/// Read an SSE stream from `reader`, write extracted content to `writer`,
/// and enforce a `limit`-byte cap on raw stream bytes consumed.
///
/// Returns an error if the stream exceeds `limit`, if a read error occurs,
/// or if writing to `writer` fails. Stops cleanly on a `data: [DONE]` line.
/// Non-`data:` lines (comments, event names, blanks) are silently skipped.
fn process_sse_stream<R: BufRead, W: Write>(
    reader: R,
    writer: &mut W,
    limit: usize,
) -> Result<(), String> {
    let mut total_bytes: usize = 0;
    for line in reader.lines() {
        let line = line.map_err(|e| format!("error reading LLM response stream: {e}"))?;
        // Count raw line bytes before parsing to bound total stream consumption
        // for malformed/adversarial responses that carry no content field.
        total_bytes += line.len();
        if total_bytes > limit {
            return Err(format!(
                "LLM response exceeded {} MB limit; aborting",
                MAX_STREAM_BYTES / 1024 / 1024
            ));
        }
        let Some(data) = line.strip_prefix("data: ") else {
            continue;
        };
        if data.trim() == "[DONE]" {
            break;
        }
        if let Ok(val) = serde_json::from_str::<serde_json::Value>(data) {
            if let Some(content) = val["choices"][0]["delta"]["content"].as_str() {
                let safe = strip_terminal_escapes(content);
                writer
                    .write_all(safe.as_bytes())
                    .map_err(|e| format!("failed to write LLM response: {e}"))?;
                writer
                    .flush()
                    .map_err(|e| format!("failed to flush stdout: {e}"))?;
            }
        }
    }
    Ok(())
}

/// POST a prompt to an OpenAI-compatible `/v1/chat/completions` endpoint
/// and stream the response to stdout.
///
/// `key` may be any non-empty string for local models (Ollama, LM Studio).
/// The endpoint must support streaming (`stream: true`).
pub(crate) fn send_prompt(
    endpoint: &str,
    model: &str,
    key: &str,
    prompt: &str,
) -> Result<(), String> {
    let url = format!("{}/chat/completions", endpoint.trim_end_matches('/'));

    let body = serde_json::json!({
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "stream": true
    });

    let agent = ureq::AgentBuilder::new()
        .timeout_connect(Duration::from_secs(10))
        .timeout_read(Duration::from_secs(300))
        .build();

    let response = agent
        .post(&url)
        .set("Authorization", &format!("Bearer {key}"))
        .set("Content-Type", "application/json")
        .send_json(body)
        .map_err(|e| match e {
            ureq::Error::Status(code, resp) => {
                let mut buf = Vec::new();
                let _ = resp
                    .into_reader()
                    .take(MAX_ERROR_BODY_BYTES)
                    .read_to_end(&mut buf);
                let body = String::from_utf8_lossy(&buf);
                format!("LLM endpoint returned HTTP {code}: {body}")
            }
            ureq::Error::Transport(t) => {
                format!("failed to reach LLM endpoint: {t}")
            }
        })?;

    let reader = BufReader::new(response.into_reader());
    let stdout = std::io::stdout();
    let mut out = stdout.lock();
    process_sse_stream(reader, &mut out, MAX_STREAM_BYTES)?;
    drop(out);
    println!();
    Ok(())
}

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

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

    fn sse_content_line(content: &str) -> String {
        let json = serde_json::json!({"choices":[{"delta":{"content": content}}]});
        format!("data: {json}\n")
    }

    // ---- strip_terminal_escapes ----

    #[test]
    fn strip_esc_removes_all_esc_bytes() {
        assert_eq!(strip_terminal_escapes("\x1b[31mred\x1b[0m"), "[31mred[0m");
        assert_eq!(strip_terminal_escapes("safe text"), "safe text");
        assert_eq!(strip_terminal_escapes("\x1b"), "");
        assert_eq!(strip_terminal_escapes("a\x1bb\x1bc"), "abc");
    }

    #[test]
    fn strip_esc_empty_string() {
        assert_eq!(strip_terminal_escapes(""), "");
    }

    // ---- process_sse_stream ----

    #[test]
    fn sse_extracts_content_from_stream() {
        let input = format!(
            "{}{}data: [DONE]\n",
            sse_content_line("hello "),
            sse_content_line("world"),
        );
        let mut out = Vec::new();
        process_sse_stream(Cursor::new(input.as_bytes()), &mut out, MAX_STREAM_BYTES).unwrap();
        assert_eq!(String::from_utf8(out).unwrap(), "hello world");
    }

    #[test]
    fn sse_strips_esc_in_content() {
        // ESC appears as  in JSON-encoded content.
        let input = format!("{}data: [DONE]\n", sse_content_line("\x1b[31mred\x1b[0m"));
        let mut out = Vec::new();
        process_sse_stream(Cursor::new(input.as_bytes()), &mut out, MAX_STREAM_BYTES).unwrap();
        let s = String::from_utf8(out).unwrap();
        assert!(
            !s.contains('\x1b'),
            "ESC bytes must be stripped; got: {s:?}"
        );
        assert!(
            s.contains("[31mred[0m"),
            "non-ESC chars must be preserved; got: {s:?}"
        );
    }

    #[test]
    fn sse_stops_at_done() {
        let input = format!(
            "{}data: [DONE]\n{}",
            sse_content_line("first"),
            sse_content_line("should-not-appear"),
        );
        let mut out = Vec::new();
        process_sse_stream(Cursor::new(input.as_bytes()), &mut out, MAX_STREAM_BYTES).unwrap();
        assert_eq!(String::from_utf8(out).unwrap(), "first");
    }

    #[test]
    fn sse_skips_non_data_lines() {
        let input = format!(
            "event: ping\n: comment\n\n{}data: [DONE]\n",
            sse_content_line("ok"),
        );
        let mut out = Vec::new();
        process_sse_stream(Cursor::new(input.as_bytes()), &mut out, MAX_STREAM_BYTES).unwrap();
        assert_eq!(String::from_utf8(out).unwrap(), "ok");
    }

    #[test]
    fn sse_returns_error_when_stream_exceeds_limit() {
        // Use a small limit so the test doesn't need to allocate 10 MB.
        let limit = 500usize;
        // Each filler line is ~100 bytes; 6 lines puts us over 500.
        let filler = format!("non-data-line: {}\n", "x".repeat(90));
        let input = filler.repeat(7);
        let mut out = Vec::new();
        let err = process_sse_stream(Cursor::new(input.as_bytes()), &mut out, limit).unwrap_err();
        assert!(
            err.contains("exceeded"),
            "error must mention exceeded; got: {err}"
        );
    }
}