//! Integration tests for StreamChunk EOF semantics.
//!
//! Verifies that the parser correctly handles 4 distinct stream termination scenarios:
//!
//! 1. **Clean stop**: `finish_reason="stop"` + `[DONE]` → yield last action (content complete)
//! 2. **Length truncated**: `finish_reason="length"` + `[DONE]` → DO NOT yield last action (content may be cut)
//! 3. **Abnormal close**: No `[DONE]` before connection drop → DO NOT yield last action (corrupted)
//! 4. **HTTP error**: Non-200 response → `StreamChunk::Error`
//!
//! Uses a TcpListener-based mock SSE server (no external HTTP mocking deps)
//! to simulate OpenAI-compatible streaming responses.

use mad_hatter::{InferError, InferErrorCode, LlmChannel, OpenAiChannel, StreamChunk};
use mad_hatter::llm::{Reflectable, Promptable, Parseable};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::task::JoinHandle;

// ─── Test types ─────────────────────────────────────────────────────────

struct TestReq;
impl Reflectable for TestReq {}
impl Promptable for TestReq {
    fn to_markdown_depth(&self, _depth: usize) -> String {
        "test".to_string()
    }
    fn schema_markdown(token: &str) -> String where Self: Sized {
        format!("TestAction-{}\nreply\ncontent-{}\n<value>\n", token, token)
    }
}

#[derive(Debug, PartialEq)]
struct TestAction {
    content: String,
}

impl Reflectable for TestAction {}
impl Promptable for TestAction {
    fn to_markdown_depth(&self, _depth: usize) -> String {
        format!("content: {}", self.content)
    }
}

impl Parseable for TestAction {
    fn type_name() -> &'static str { "TestAction" }

    fn from_markdown(text: &str, token: &str) -> Result<Vec<Self>, String>
    where Self: Sized {
        let sep = format!("TestAction-{}", token);
        let content_tag = format!("content-{}", token);
        let mut actions = Vec::new();
        // Split by separator lines
        let parts: Vec<&str> = text.split(&format!("{}\n", sep)).collect();
        for part in parts.iter().skip(1) { // skip text before first separator
            // Expect: "reply\ncontent-{token}\n<value>\n..."
            // Find "content-{token}\n" marker
            if let Some(content_start) = part.find(&format!("{}\n", content_tag)) {
                let after = content_start + content_tag.len() + 1;
                let remaining = &part[after..];
                // Content is everything until next TestAction- or end
                let content_end = remaining.find("TestAction-").unwrap_or(remaining.len());
                let content = remaining[..content_end].trim();
                if !content.is_empty() {
                    actions.push(TestAction { content: content.to_string() });
                }
            }
        }
        if actions.is_empty() {
            return Err(format!("No valid TestAction found in: {}", text));
        }
        Ok(actions)
    }
}

// ─── Mock SSE server ────────────────────────────────────────────────────

/// Spec for mock server behavior
struct SseScript {
    /// Sequence of SSE data events to emit (without "data: " prefix, without trailing \n\n)
    events: Vec<String>,
    /// Whether to send `data: [DONE]` at the end
    send_done: bool,
    /// Whether to close connection abruptly mid-stream (after N events)
    abrupt_close_after: Option<usize>,
    /// HTTP status code (default 200)
    status: u16,
    /// Error body for non-200 responses
    error_body: Option<String>,
}

impl SseScript {
    fn clean_stop(content_chunks: Vec<&str>) -> Self {
        let mut events: Vec<String> = content_chunks
            .iter()
            .map(|c| sse_content_delta(c))
            .collect();
        // Final chunk with finish_reason=stop
        events.push(sse_finish("stop"));
        Self {
            events,
            send_done: true,
            abrupt_close_after: None,
            status: 200,
            error_body: None,
        }
    }

    fn length_truncated(content_chunks: Vec<&str>) -> Self {
        let mut events: Vec<String> = content_chunks
            .iter()
            .map(|c| sse_content_delta(c))
            .collect();
        events.push(sse_finish("length"));
        Self {
            events,
            send_done: true,
            abrupt_close_after: None,
            status: 200,
            error_body: None,
        }
    }

    fn abrupt_close(content_chunks: Vec<&str>) -> Self {
        let events: Vec<String> = content_chunks
            .iter()
            .map(|c| sse_content_delta(c))
            .collect();
        Self {
            events,
            send_done: false,
            abrupt_close_after: None, // close after sending all events
            status: 200,
            error_body: None,
        }
    }

    fn http_error(status: u16, body: &str) -> Self {
        Self {
            events: vec![],
            send_done: false,
            abrupt_close_after: None,
            status,
            error_body: Some(body.to_string()),
        }
    }
}

fn sse_content_delta(content: &str) -> String {
    let escaped = content.replace('\\', "\\\\").replace('"', "\\\"").replace('\n', "\\n");
    format!(
        r#"{{"choices":[{{"delta":{{"content":"{}"}}, "finish_reason":null}}]}}"#,
        escaped
    )
}

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

async fn start_mock_sse_server(script: SseScript) -> (u16, JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let handle = tokio::spawn(async move {
        loop {
            let (mut stream, _) = match listener.accept().await {
                Ok(v) => v,
                Err(_) => break,
            };
            let events = script.events.clone();
            let send_done = script.send_done;
            let status = script.status;
            let error_body = script.error_body.clone();
            tokio::spawn(async move {
                // Read & discard request headers+body (minimal parse)
                let mut buf = vec![0u8; 16384];
                let mut total = 0usize;
                let _ = tokio::time::timeout(
                    std::time::Duration::from_millis(500),
                    async {
                        loop {
                            let n = match stream.read(&mut buf[total..]).await {
                                Ok(0) => break,
                                Ok(n) => n,
                                Err(_) => return,
                            };
                            total += n;
                            if total >= 4 && buf[..total].windows(4).any(|w| w == b"\r\n\r\n") {
                                break;
                            }
                        }
                    },
                )
                .await;

                // Write response
                if status != 200 {
                    let body = error_body.unwrap_or_else(|| "error".to_string());
                    let resp = format!(
                        "HTTP/1.1 {} Error\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
                        status,
                        body.len(),
                        body
                    );
                    let _ = stream.write_all(resp.as_bytes()).await;
                    let _ = stream.shutdown().await;
                    return;
                }

                // 200 OK SSE stream
                let head = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n";
                let _ = stream.write_all(head.as_bytes()).await;

                for event in &events {
                    let payload = format!("data: {}\n\n", event);
                    let chunk = format!("{:x}\r\n{}\r\n", payload.len(), payload);
                    if stream.write_all(chunk.as_bytes()).await.is_err() {
                        return;
                    }
                    // Small delay between events
                    tokio::time::sleep(std::time::Duration::from_millis(5)).await;
                }

                if send_done {
                    let payload = "data: [DONE]\n\n";
                    let chunk = format!("{:x}\r\n{}\r\n", payload.len(), payload);
                    let _ = stream.write_all(chunk.as_bytes()).await;
                }

                // End chunked stream (or don't, for abrupt close)
                if send_done {
                    let _ = stream.write_all(b"0\r\n\r\n").await;
                    let _ = stream.shutdown().await;
                } else {
                    // Abrupt close: just drop the stream without proper chunked end
                    drop(stream);
                }
            });
        }
    });
    (port, handle)
}

// ─── Tests ──────────────────────────────────────────────────────────────

/// Scenario 1: Clean stop — finish_reason="stop" + [DONE] → yield last action
#[tokio::test]
async fn test_clean_stop_yields_last_action() {
    // Extract token from the request prompt — the engine generates a random token per-call.
    // Our mock server ignores the prompt and sends a response crafted for a known token.
    // We need to use a channel that lets us see the prompt first to extract the token.
    // Since OpenAiChannel can't do this, we build the SSE content to match any token by
    // emitting a fixed-format response and the parser uses the actual token in its state.

    // Alternative: use infer() directly but control the token via the prompt inspection.
    // For simplicity here, we skip the round-trip and verify via StreamInfer directly
    // with a synthetic channel... but the goal is end-to-end, so let's use OpenAiChannel.

    // Workaround: the token is embedded in the schema which is sent in the prompt.
    // We'll extract it by instrumenting... or use a known-fixed prompt with a deterministic token.

    // Simpler: verify at the channel level (StreamChunk emission) without going through
    // the full infer() pipeline. This directly tests the EOF semantics in llm.rs.

    let script = SseScript::clean_stop(vec![
        "TestAction-TOK\n",
        "reply\n",
        "content-TOK\n",
        "hello world",
    ]);
    let (port, _handle) = start_mock_sse_server(script).await;

    let channel = OpenAiChannel::new(
        format!("http://127.0.0.1:{}/v1/chat/completions", port),
        "gpt-4".to_string(),
        "fake-key".to_string(),
    );

    let mut rx = channel.start_stream("prompt".to_string()).unwrap();
    let mut chunks = Vec::new();
    while let Some(chunk) = rx.recv().await {
        chunks.push(chunk);
    }

    // Verify we got: Content chunks + Eof { clean: true }
    assert!(!chunks.is_empty(), "should receive at least one chunk");
    let last = chunks.last().unwrap();
    match last {
        StreamChunk::Eof { clean: true } => {}
        other => panic!("Expected Eof {{ clean: true }}, got {:?}", match other {
            StreamChunk::Content(s) => format!("Content({:?})", s),
            StreamChunk::Usage(_) => "Usage".to_string(),
            StreamChunk::Error(e) => format!("Error({:?})", e),
            StreamChunk::Eof { clean } => format!("Eof {{ clean: {} }}", clean),
        }),
    }

    // Verify content was transmitted
    let concat: String = chunks.iter().filter_map(|c| match c {
        StreamChunk::Content(s) => Some(s.as_str()),
        _ => None,
    }).collect();
    assert!(concat.contains("TestAction-TOK"), "content should contain action start");
    assert!(concat.contains("hello world"), "content should contain body");
}

/// Scenario 2: Length truncated — finish_reason="length" + [DONE] → Eof { clean: false }
#[tokio::test]
async fn test_length_truncated_signals_dirty_eof() {
    let script = SseScript::length_truncated(vec![
        "TestAction-TOK\n",
        "reply\n",
        "content-TOK\n",
        "partial con", // truncated mid-word
    ]);
    let (port, _handle) = start_mock_sse_server(script).await;

    let channel = OpenAiChannel::new(
        format!("http://127.0.0.1:{}/v1/chat/completions", port),
        "gpt-4".to_string(),
        "fake-key".to_string(),
    );

    let mut rx = channel.start_stream("prompt".to_string()).unwrap();
    let mut chunks = Vec::new();
    while let Some(chunk) = rx.recv().await {
        chunks.push(chunk);
    }

    // Verify last chunk is Eof { clean: false }
    let last = chunks.last().unwrap();
    match last {
        StreamChunk::Eof { clean: false } => {}
        other => panic!("Expected Eof {{ clean: false }} (length truncation), got {:?}", match other {
            StreamChunk::Content(s) => format!("Content({:?})", s),
            StreamChunk::Usage(_) => "Usage".to_string(),
            StreamChunk::Error(e) => format!("Error({:?})", e),
            StreamChunk::Eof { clean } => format!("Eof {{ clean: {} }}", clean),
        }),
    }
}

/// Scenario 3: Abrupt close — connection drops without [DONE] → no Eof signal at all
#[tokio::test]
async fn test_abrupt_close_no_eof_signal() {
    let script = SseScript::abrupt_close(vec![
        "TestAction-TOK\n",
        "reply\n",
        "content-TOK\n",
        "half",
    ]);
    let (port, _handle) = start_mock_sse_server(script).await;

    let channel = OpenAiChannel::new(
        format!("http://127.0.0.1:{}/v1/chat/completions", port),
        "gpt-4".to_string(),
        "fake-key".to_string(),
    );

    let mut rx = channel.start_stream("prompt".to_string()).unwrap();
    let mut chunks = Vec::new();
    while let Some(chunk) = rx.recv().await {
        chunks.push(chunk);
    }

    // With abrupt close, there should be NO Eof variant at all
    // (spawn_sse_reader only sends Eof on [DONE] path; abrupt break in stream just closes channel)
    let has_eof = chunks.iter().any(|c| matches!(c, StreamChunk::Eof { .. }));
    assert!(!has_eof, "abrupt close should NOT emit Eof; chunks: {:?}", chunks.iter().map(|c| match c {
        StreamChunk::Content(s) => format!("Content({})", s.len()),
        StreamChunk::Usage(_) => "Usage".to_string(),
        StreamChunk::Error(e) => format!("Error({})", e),
        StreamChunk::Eof { clean } => format!("Eof(clean={})", clean),
    }).collect::<Vec<_>>());
}

/// Scenario 4: HTTP error — non-200 response → StreamChunk::Error
#[tokio::test]
async fn test_http_error_emits_error_chunk() {
    let script = SseScript::http_error(500, r#"{"error":"internal server error"}"#);
    let (port, _handle) = start_mock_sse_server(script).await;

    let channel = OpenAiChannel::new(
        format!("http://127.0.0.1:{}/v1/chat/completions", port),
        "gpt-4".to_string(),
        "fake-key".to_string(),
    );

    let mut rx = channel.start_stream("prompt".to_string()).unwrap();
    let mut chunks = Vec::new();
    while let Some(chunk) = rx.recv().await {
        chunks.push(chunk);
    }

    // Should have exactly one Error chunk
    assert_eq!(chunks.len(), 1, "HTTP error should emit exactly one chunk; got: {}", chunks.len());
    match &chunks[0] {
        StreamChunk::Error(msg) => {
            assert!(msg.contains("500"), "error should mention HTTP 500: {}", msg);
            assert!(msg.contains("internal server error"), "error should include body: {}", msg);
        }
        other => panic!("Expected Error, got {:?}", match other {
            StreamChunk::Content(s) => format!("Content({:?})", s),
            StreamChunk::Usage(_) => "Usage".to_string(),
            StreamChunk::Error(e) => format!("Error({:?})", e),
            StreamChunk::Eof { clean } => format!("Eof {{ clean: {} }}", clean),
        }),
    }
}

// ─── Note on infer() integration ────────────────────────────────────────
//
// Full infer() pipeline behavior with EOF semantics is covered by test_infer.rs:
//   - test_infer_single_action (clean EOF + single action → yield works)
//   - test_infer_multiple_actions (multiple separators delimit actions)
//   - test_infer_channel_error (Error chunk surfaces correctly)
//   - ... and 9 more integration tests
//
// This file focuses on the channel-level StreamChunk emission contract,
// using a real TcpListener mock SSE server to exercise the spawn_sse_reader
// path end-to-end (HTTP parsing → SseChoice deserialization → StreamChunk
// variant selection). The 4 scenarios above prove:
//   1. finish_reason="stop" + [DONE] → StreamChunk::Eof { clean: true }
//   2. finish_reason="length" + [DONE] → StreamChunk::Eof { clean: false }
//   3. No [DONE] + connection drop → channel closes without Eof variant
//   4. HTTP non-200 → StreamChunk::Error with status + body
//
// StreamInfer::next() consumes these StreamChunk variants and enforces the
// yield discipline: last unterminated block is only yielded on clean EOF.
// That discipline is exercised by test_infer.rs.