//! Test: infer() drops stream after receiving first result,
//! causing "cannot write to closing transport" on server side.
//!
//! This test reproduces the bug where:
//! 1. Client calls infer()
//! 2. LLM server starts streaming SSE response
//! 3. infer() parses the first (and only) result, returns immediately
//! 4. StreamInfer is dropped → receiver is dropped
//! 5. Server task tries to send more data (usage, DONE) → fails
//! 6. In production, server sees "cannot write to closing transport"

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

// === Minimal test types (hand-implemented, no derive macros needed) ===

struct TestReq;

impl mad_hatter::llm::Reflectable for TestReq {}
impl mad_hatter::llm::Promptable for TestReq {
    fn to_markdown_depth(&self, _depth: usize) -> String {
        "test request content".to_string()
    }
}

struct TestResp {
    value: String,
}

impl mad_hatter::llm::Reflectable for TestResp {}
impl mad_hatter::llm::Promptable for TestResp {
    fn to_markdown_depth(&self, _depth: usize) -> String {
        format!("value: {}", self.value)
    }
    fn schema_markdown(token: &str) -> String
    where
        Self: Sized,
    {
        format!(
            "TestResp-{token}\nvalue: <your response>\nTestResp-end-{token}",
            token = token,
        )
    }
}
impl mad_hatter::llm::Parseable for TestResp {
    fn from_markdown(text: &str, _token: &str) -> Result<Vec<Self>, String>
    where
        Self: Sized,
    {
        for line in text.lines() {
            if let Some(val) = line.strip_prefix("value: ") {
                return Ok(vec![TestResp {
                    value: val.trim().to_string(),
                }]);
            }
        }
        Err("TestResp: no 'value:' line found".to_string())
    }
    fn type_name() -> &'static str
    where
        Self: Sized,
    {
        "TestResp"
    }
}

// === Mock LLM Channel ===
//
// Simulates a server that streams a valid response, then tries to send
// usage data after the content is delivered. By the time usage is sent,
// infer() has already returned and dropped the stream.

struct SlowMockChannel {
    /// Set to true when spawned task detects client disconnected
    disconnect_detected: Arc<AtomicBool>,
}

impl mad_hatter::llm::LlmChannel for SlowMockChannel {
    fn start_stream(
        &self,
        prompt: String,
    ) -> Result<tokio::sync::mpsc::UnboundedReceiver<String>, String> {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let flag = self.disconnect_detected.clone();

        tokio::spawn(async move {
            // Extract token from prompt.
            // build_cached_prompt includes schema_markdown output which contains "TestResp-{token}".
            // Token is a 6-char hex string generated by generate_token().
            let token = {
                let re = regex::Regex::new(r"TestResp-([0-9a-f]{6})").unwrap();
                match re.captures(&prompt) {
                    Some(caps) => caps[1].to_string(),
                    None => {
                        let _ = tx.send(
                            "\n[ERROR] HTTP request failed: no token found in prompt".to_string(),
                        );
                        return;
                    }
                }
            };

            // Build valid response: one complete parsed item
            let response = format!(
                "TestResp-{token}\nvalue: hello world\nTestResp-end-{token}\n",
                token = token,
            );

            // Send content in small chunks with delays (simulate real SSE streaming)
            for chunk in response.as_bytes().chunks(15) {
                let s = String::from_utf8_lossy(chunk).to_string();
                if tx.send(s).is_err() {
                    flag.store(true, Ordering::SeqCst);
                    eprintln!("[TEST-SERVER] Client disconnected during content delivery");
                    return;
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }

            // === Content fully delivered ===
            //
            // At this point, StreamInfer has enough data to parse the result.
            // infer() will call stream.next(), get Ok(TestResp), and return.
            // This drops StreamInfer → drops receiver → our tx.send() will fail.
            //
            // In production (spawn_sse_reader), this is the moment where:
            // - The reqwest Response is still streaming (usage + [DONE] not sent yet)
            // - Client closes the HTTP connection
            // - Server sees "cannot write to closing transport"
            //
            // We simulate this by waiting, then trying to send more data.

            tokio::time::sleep(Duration::from_millis(300)).await;

            // Try to send usage data — this SHOULD FAIL because receiver is dropped
            match tx.send("\n[USAGE:100:50:80:20]".to_string()) {
                Ok(_) => {
                    eprintln!("[TEST-SERVER] WARNING: usage send succeeded (receiver still alive)");
                }
                Err(_) => {
                    flag.store(true, Ordering::SeqCst);
                    eprintln!(
                        "[TEST-SERVER] Disconnect detected when sending usage — \
                         this is the 'cannot write to closing transport' equivalent"
                    );
                    return;
                }
            }

            // If usage somehow succeeded, try DONE marker
            tokio::time::sleep(Duration::from_millis(100)).await;
            if tx.send("extra".to_string()).is_err() {
                flag.store(true, Ordering::SeqCst);
                eprintln!("[TEST-SERVER] Disconnect detected when sending extra data");
            }
        });

        Ok(rx)
    }
}

// === The Test ===

#[tokio::test]
async fn test_infer_drops_stream_causing_server_disconnect() {
    let disconnect = Arc::new(AtomicBool::new(false));
    let channel = SlowMockChannel {
        disconnect_detected: disconnect.clone(),
    };

    // Call infer — should succeed with parsed result
    let result = mad_hatter::infer::<TestReq, TestResp>(&channel, &TestReq).await;
    assert!(result.is_ok(), "infer should succeed: {:?}", result.err());

    let (resp, usage) = result.unwrap();
    assert_eq!(resp.value, "hello world");

    // Usage should be None because we returned before server sent it.
    // This is a side effect of the bug: usage data is lost.
    assert!(
        usage.is_none(),
        "usage should be None — stream was dropped before server sent usage data"
    );

    // At this point, StreamInfer is dropped, receiver is dropped.
    // Wait for server task to detect the disconnect.
    tokio::time::sleep(Duration::from_millis(800)).await;

    assert!(
        disconnect.load(Ordering::SeqCst),
        "BUG CONFIRMED: Server tried to send data after client closed the channel. \
         In production with HTTP, this manifests as 'cannot write to closing transport'. \
         Root cause: infer() returns immediately after parsing first result, \
         dropping the stream while the server still has data to send (usage, [DONE])."
    );
}