use mad_hatter::{Promptable, Parseable, LlmChannel, stream_infer, stream_infer_with_on_text, InferError};
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc::UnboundedReceiver;

// --- Test types ---

#[derive(Promptable)]
#[prompt("Test request")]
struct StreamReq {
    #[prompt("prompt")]
    prompt: String,
}

#[derive(Parseable, PartialEq, Debug)]
enum StreamAction {
    #[prompt("思考")]
    Think {
        #[prompt("内容")]
        content: String,
    },
    #[prompt("回复")]
    Reply {
        #[prompt("内容")]
        content: String,
    },
    #[prompt("等待")]
    Idle,
}

// Helper: extract token from prompt
fn extract_token(prompt: &str, type_name: &str) -> String {
    let marker = format!("{}-", type_name);
    for line in prompt.lines() {
        if let Some(pos) = line.find(&marker) {
            let after = &line[pos + marker.len()..];
            // Token is hex chars until non-hex
            let token: String = after.chars().take_while(|c| c.is_ascii_hexdigit()).collect();
            if !token.is_empty() {
                return token;
            }
        }
    }
    panic!("Token not found in prompt for type {}", type_name);
}

// Helper: create a mock channel that sends chunks via unbounded_channel
fn mock_channel_rx(chunks: Vec<String>) -> UnboundedReceiver<String> {
    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
    for chunk in chunks {
        tx.send(chunk).ok();
    }
    // tx dropped here → channel closed
    rx
}

// --- Tests ---

#[tokio::test]
async fn test_stream_infer_yields_elements_one_by_one() {
    struct ChunkedChannel;

    impl LlmChannel for ChunkedChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let chunks = vec![
                format!("StreamAction-{}\n", token),
                "think\n".to_string(),
                format!("content-{}\n", token),
                "正在思考...\n".to_string(),
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "回答完毕。\n".to_string(),
                format!("StreamAction-{}\n", token),
                "idle\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(&ChunkedChannel, &request, true).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 3);
    assert_eq!(results[0], StreamAction::Think { content: "正在思考...".to_string() });
    assert_eq!(results[1], StreamAction::Reply { content: "回答完毕。".to_string() });
    assert_eq!(results[2], StreamAction::Idle);
}

#[tokio::test]
async fn test_stream_infer_single_chunk() {
    // All data arrives in one chunk (like mock channel)
    struct OneChunkChannel;

    impl LlmChannel for OneChunkChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let response = format!(
                "StreamAction-{t}\nthink\ncontent-{t}\n深度思考中\nStreamAction-{t}\nreply\ncontent-{t}\n最终答案\nStreamAction-end-{t}",
                t = token
            );
            Ok(mock_channel_rx(vec![response]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(&OneChunkChannel, &request, true).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 2);
    assert_eq!(results[0], StreamAction::Think { content: "深度思考中".to_string() });
    assert_eq!(results[1], StreamAction::Reply { content: "最终答案".to_string() });
}

#[tokio::test]
async fn test_stream_infer_missing_end_marker() {
    struct TruncatedChannel;

    impl LlmChannel for TruncatedChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let response = format!(
                "StreamAction-{t}\nthink\ncontent-{t}\n被截断了",
                t = token
            );
            Ok(mock_channel_rx(vec![response]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(&TruncatedChannel, &request, true).await.unwrap();

    // First next() should return the truncation error
    let result = stream.next().await;
    assert!(result.is_some());
    let err = result.unwrap();
    assert!(err.is_err());
    let err_msg = err.unwrap_err();
    assert!(err_msg.message.contains("No valid element found") || err_msg.message.contains("end marker"), "Error should mention missing element or end marker: {}", err_msg);
}

#[tokio::test]
async fn test_stream_infer_channel_error() {
    struct FailChannel;

    impl LlmChannel for FailChannel {
        fn start_stream(&self, _prompt: String) -> Result<UnboundedReceiver<String>, String> {
            Err("Connection refused".to_string())
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let result = stream_infer::<StreamReq, StreamAction>(&FailChannel, &request, true).await;
    assert!(result.is_err());
    match result {
        Err(e) => assert!(e.message.contains("Connection refused"), "unexpected error: {}", e),
        Ok(_) => panic!("expected error but got Ok"),
    }
}

#[tokio::test]
async fn test_stream_infer_token_accessible() {
    struct TokenChannel;

    impl LlmChannel for TokenChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let response = format!("StreamAction-{t}\nidle\nStreamAction-end-{t}", t = token);
            Ok(mock_channel_rx(vec![response]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let stream = stream_infer::<StreamReq, StreamAction>(&TokenChannel, &request, true).await.unwrap();

    // Token should be accessible
    let token = stream.token().to_string();
    assert!(!token.is_empty());
    assert!(token.chars().all(|c| c.is_ascii_hexdigit()));
}

#[tokio::test]
async fn test_stream_infer_gradual_chunks() {
    // Simulate very small chunks (like real SSE streaming)
    struct GradualChannel;

    impl LlmChannel for GradualChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let full = format!(
                "StreamAction-{t}\nreply\ncontent-{t}\nHello World\nStreamAction-end-{t}",
                t = token
            );
            // Split into 2-char chunks
            let chunks: Vec<String> = full.as_bytes()
                .chunks(2)
                .map(|c| String::from_utf8_lossy(c).to_string())
                .collect();
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(&GradualChannel, &request, true).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Reply { content: "Hello World".to_string() });
}

#[tokio::test]
async fn test_stream_infer_empty_response() {
    // Only separator + end marker, no actual elements
    struct EmptyChannel;

    impl LlmChannel for EmptyChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let response = format!("StreamAction-end-{t}", t = token);
            Ok(mock_channel_rx(vec![response]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(&EmptyChannel, &request, true).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    assert_eq!(results.len(), 0);
}

#[tokio::test]
async fn test_stream_infer_with_code_block_wrapper() {
    // LLM wraps output in ```
    struct CodeBlockChannel;

    impl LlmChannel for CodeBlockChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let response = format!(
                "```\nStreamAction-{t}\nreply\ncontent-{t}\n代码块内的回复\nStreamAction-end-{t}\n```",
                t = token
            );
            Ok(mock_channel_rx(vec![response]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(&CodeBlockChannel, &request, true).await.unwrap();
    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    // The ``` lines don't match separator format, so they're ignored as noise
    assert_eq!(results.len(), 1);
    assert!(results[0].is_ok());
    assert_eq!(results[0].as_ref().unwrap(), &StreamAction::Reply { content: "代码块内的回复".to_string() });
}

#[tokio::test]
async fn test_stream_infer_on_text_receives_all_chunks() {
    // Channel that returns response in multiple chunks
    struct ChunkedCallbackChannel;

    impl LlmChannel for ChunkedCallbackChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let chunks = vec![
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "回调测试内容\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let collected = Arc::new(Mutex::new(Vec::<String>::new()));
    let collected_clone = collected.clone();

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &ChunkedCallbackChannel,
        &request,
        Some(Box::new(move |chunk: &str| {
            collected_clone.lock().unwrap().push(chunk.to_string());
        })),
        None,
        None,
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }

    // Verify parsing works
    assert_eq!(results.len(), 1);
    assert!(results[0].is_ok());
    assert_eq!(results[0].as_ref().unwrap(), &StreamAction::Reply { content: "回调测试内容".to_string() });

    // Verify callback received all 4 chunks
    let chunks = collected.lock().unwrap();
    assert_eq!(chunks.len(), 5);
    assert!(chunks[0].starts_with("StreamAction-"));
    assert_eq!(chunks[1], "reply\n");
    assert!(chunks[2].starts_with("content-"));
    assert_eq!(chunks[3], "回调测试内容\n");
    assert!(chunks[4].starts_with("StreamAction-end-"));
}

#[tokio::test]
async fn test_stream_infer_on_text_none_equivalent() {
    // Same as regular stream_infer
    struct SimpleCallbackChannel;

    impl LlmChannel for SimpleCallbackChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let response = format!(
                "StreamAction-{t}\nreply\ncontent-{t}\n无回调\nStreamAction-end-{t}\n",
                t = token
            );
            Ok(mock_channel_rx(vec![response]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &SimpleCallbackChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    assert_eq!(results.len(), 1);
    assert!(results[0].is_ok());
    assert_eq!(results[0].as_ref().unwrap(), &StreamAction::Reply { content: "无回调".to_string() });
}

// === preamble rejection tests ===

#[tokio::test]
async fn test_stream_infer_preamble_returns_error() {
    // Preamble before first action should cause an error (zero tolerance)
    struct PreambleChannel;

    impl LlmChannel for PreambleChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("I need to think about this carefully.\nLet me analyze the situation.\nStreamAction-{}\nidle\nStreamAction-end-{}\n", token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &PreambleChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let result = stream.next().await;
    assert!(result.is_some(), "should return an error result");
    let err = result.unwrap();
    assert!(err.is_err(), "preamble should cause error");
    assert!(err.unwrap_err().message.contains("FORMAT VIOLATION"), "error should mention unexpected content");

    // Stream should be done after preamble error
    assert!(stream.next().await.is_none(), "stream should be done after preamble error");
}

#[tokio::test]
async fn test_stream_infer_no_preamble_works_normally() {
    // When there's no preamble, everything works normally
    struct NoPreambleChannel;

    impl LlmChannel for NoPreambleChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("StreamAction-{}\nidle\nStreamAction-end-{}\n", token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &NoPreambleChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    assert_eq!(results.len(), 1);
    assert!(results[0].is_ok());
}

#[tokio::test]
async fn test_stream_infer_whitespace_only_preamble_ok() {
    // Whitespace-only content before first separator should be silently ignored (not an error)
    struct WhitespacePreambleChannel;

    impl LlmChannel for WhitespacePreambleChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("\n  \n```\nStreamAction-{}\nidle\nStreamAction-end-{}\n", token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &WhitespacePreambleChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    assert_eq!(results.len(), 1, "whitespace/backtick-only preamble should be ignored");
    assert!(results[0].is_ok());
}

// === Cancel tests ===

#[tokio::test]
async fn test_stream_infer_cancel_mid_stream() {
    use std::sync::atomic::{AtomicBool, Ordering};

    struct SlowChannel;
    impl LlmChannel for SlowChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let chunks = vec![
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "first action\n".to_string(),
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "second action\n".to_string(),
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "third action\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let cancel = Arc::new(AtomicBool::new(false));
    let cancel_clone = cancel.clone();

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &SlowChannel,
        &request,
        None,
        None,
        Some(cancel_clone),
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(item) = stream.next().await {
        results.push(item.unwrap());
        // Cancel after first action
        cancel.store(true, Ordering::Relaxed);
    }

    // Should have gotten first action then stopped
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Reply { content: "first action".to_string() });
}

#[tokio::test]
async fn test_stream_infer_cancel_none_no_effect() {
    // Verify that cancel=None doesn't affect normal operation
    let request = StreamReq { prompt: "test".to_string() };

    struct NormalChannel;
    impl LlmChannel for NormalChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("StreamAction-{}\nreply\ncontent-{}\nworks\nStreamAction-end-{}\n", token, token, token),
            ]))
        }
    }

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &NormalChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0].as_ref().unwrap(), &StreamAction::Reply { content: "works".to_string() });
}

#[tokio::test]
async fn test_stream_infer_cancel_before_start() {
    use std::sync::atomic::AtomicBool;

    struct AnyChannel;
    impl LlmChannel for AnyChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("StreamAction-{}\nreply\ncontent-{}\nshould not see\nStreamAction-end-{}\n", token, token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let cancel = Arc::new(AtomicBool::new(true)); // Already cancelled

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &AnyChannel,
        &request,
        None,
        None,
        Some(cancel),
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result);
    }
    assert_eq!(results.len(), 0); // Nothing parsed
}

// --- Thinking tests ---

#[tokio::test]
async fn test_stream_thinking_then_action() {
    struct ThinkingChannel;

    impl LlmChannel for ThinkingChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let chunks = vec![
format!("<planning-{}>\nI need to think about this carefully.\n</planning-{}>\n", token, token),
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "hello world\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let collected_thinking = Arc::new(Mutex::new(Vec::<String>::new()));
    let thinking_clone = collected_thinking.clone();

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &ThinkingChannel,
        &request,
        None,
        None,
        None,
        Some(Box::new(move |thinking: &str| {
            thinking_clone.lock().unwrap().push(thinking.to_string());
        })), true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Reply { content: "hello world".to_string() });

    let thinking = collected_thinking.lock().unwrap();
    assert_eq!(thinking.len(), 1);
    assert!(thinking[0].contains("think about this carefully"));
}

#[tokio::test]
async fn test_stream_thinking_no_callback() {
    struct ThinkingNoCallbackChannel;

    impl LlmChannel for ThinkingNoCallbackChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let chunks = vec![
format!("<planning-{}>\nSome internal reasoning\n</planning-{}>\n", token, token),
                format!("StreamAction-{}\n", token),
                "idle\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    // on_thinking = None — thinking should be silently discarded
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &ThinkingNoCallbackChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Idle);
}

#[tokio::test]
async fn test_stream_thinking_with_leading_whitespace() {
    struct ThinkingWhitespaceChannel;

    impl LlmChannel for ThinkingWhitespaceChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            let chunks = vec![
                "\n\n".to_string(),
format!("<planning-{}>\nthinking with leading whitespace\n</planning-{}>\n", token, token),
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "result\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let collected_thinking = Arc::new(Mutex::new(Vec::<String>::new()));
    let thinking_clone = collected_thinking.clone();

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &ThinkingWhitespaceChannel,
        &request,
        None,
        None,
        None,
        Some(Box::new(move |thinking: &str| {
            thinking_clone.lock().unwrap().push(thinking.to_string());
        })), true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Reply { content: "result".to_string() });

    let thinking = collected_thinking.lock().unwrap();
    assert_eq!(thinking.len(), 1);
    assert!(thinking[0].contains("leading whitespace"));
}

#[tokio::test]
async fn test_stream_thinking_gradual_chunks() {
    struct ThinkingGradualChannel;

    impl LlmChannel for ThinkingGradualChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            // Thinking arrives in small chunks
            let planning_open = format!("<planning-{}>", token);
                let planning_close = format!("</planning-{}>", token);
                let mid = planning_open.len() / 2;
                let chunks = vec![
                planning_open[..mid].to_string(),
                format!("{}\nLine 1\n", &planning_open[mid..]),
                format!("Line 2\n{}\n", planning_close),
                format!("StreamAction-{}\n", token),
                "reply\n".to_string(),
                format!("content-{}\n", token),
                "done\n".to_string(),
                format!("StreamAction-end-{}\n", token),
            ];
            Ok(mock_channel_rx(chunks))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let collected_thinking = Arc::new(Mutex::new(Vec::<String>::new()));
    let thinking_clone = collected_thinking.clone();

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &ThinkingGradualChannel,
        &request,
        None,
        None,
        None,
        Some(Box::new(move |thinking: &str| {
            thinking_clone.lock().unwrap().push(thinking.to_string());
        })), true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Reply { content: "done".to_string() });

    let thinking = collected_thinking.lock().unwrap();
    assert_eq!(thinking.len(), 1);
    assert!(thinking[0].contains("Line 1"));
    assert!(thinking[0].contains("Line 2"));
}


// === planning state machine tests ===

#[tokio::test]
async fn test_stream_planning_unclosed_returns_truncated() {
    // Planning block opened but never closed — should be Truncated, NOT FormatViolation
    struct UnclosedPlanningChannel;

    impl LlmChannel for UnclosedPlanningChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            // Send planning open tag + some content, but NO close tag, then stream ends
            Ok(mock_channel_rx(vec![
                format!("<planning-{}>\n", token),
                "I am thinking about this problem...\n".to_string(),
                "Still thinking...\n".to_string(),
                // No </planning-{token}> — stream ends here
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &UnclosedPlanningChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let result = stream.next().await;
    assert!(result.is_some(), "should return an error result");
    let err = result.unwrap();
    assert!(err.is_err(), "unclosed planning should cause error");
    let error = err.unwrap_err();
    // Must be Truncated, NOT FormatViolation
    assert_eq!(error.code, mad_hatter::InferErrorCode::Truncated,
        "unclosed planning should be Truncated, not FormatViolation. Got: {}", error.message);
    assert!(error.message.contains("Planning block opened but never closed"),
        "error message should mention unclosed planning. Got: {}", error.message);

    // Stream should be done
    assert!(stream.next().await.is_none(), "stream should be done after truncated error");
}

#[tokio::test]
async fn test_stream_garbage_preamble_still_format_violation() {
    // Garbage text with no planning at all — should remain FormatViolation
    struct GarbagePreambleChannel;

    impl LlmChannel for GarbagePreambleChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                "Here is my analysis of the situation.\n".to_string(),
                format!("StreamAction-{}\nidle\nStreamAction-end-{}\n", token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &GarbagePreambleChannel,
        &request,
        None,
        None,
        None,
        None, true,
    ).await.unwrap();

    let result = stream.next().await;
    assert!(result.is_some(), "should return an error result");
    let err = result.unwrap();
    assert!(err.is_err(), "garbage preamble should cause error");
    let error = err.unwrap_err();
    assert_eq!(error.code, mad_hatter::InferErrorCode::FormatViolation,
        "garbage preamble should be FormatViolation. Got: {}", error.message);
    assert!(error.message.contains("FORMAT VIOLATION"),
        "error message should contain FORMAT VIOLATION. Got: {}", error.message);
}

#[tokio::test]
async fn test_stream_planning_closed_then_content_works() {
    // Planning properly closed, then normal content — should work fine
    struct ClosedPlanningChannel;

    impl LlmChannel for ClosedPlanningChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("<planning-{}>\nQuick plan: just reply.\n</planning-{}>\n", token, token),
                format!("StreamAction-{}\nreply\ncontent-{}\nHello!\nStreamAction-end-{}\n", token, token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let collected_thinking = Arc::new(Mutex::new(Vec::<String>::new()));
    let thinking_clone = collected_thinking.clone();

    let mut stream = stream_infer_with_on_text::<StreamReq, StreamAction>(
        &ClosedPlanningChannel,
        &request,
        None,
        None,
        None,
        Some(Box::new(move |thinking: &str| {
            thinking_clone.lock().unwrap().push(thinking.to_string());
        })), true,
    ).await.unwrap();

    let mut results = Vec::new();
    while let Some(result) = stream.next().await {
        results.push(result.unwrap());
    }
    assert_eq!(results.len(), 1);
    assert_eq!(results[0], StreamAction::Reply { content: "Hello!".to_string() });

    let thinking = collected_thinking.lock().unwrap();
    assert_eq!(thinking.len(), 1);
    assert!(thinking[0].contains("Quick plan"));
}

#[tokio::test]
async fn test_stream_planning_hallucination_after_action() {
    // LLM outputs an action, then emits a planning tag — hallucination
    struct HallucinationChannel;

    impl LlmChannel for HallucinationChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                // First action — normal
                format!("StreamAction-{}\nreply\ncontent-{}\nHello!\n", token, token),
                // Second separator starts, then planning tag appears (hallucination)
                format!("StreamAction-{}\n<planning-{}>\nLet me reconsider...\n</planning-{}>\n", token, token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(
        &HallucinationChannel,
        &request,
        true,
    ).await.unwrap();

    // First action should succeed
    let first = stream.next().await;
    assert!(first.is_some());
    let first_result = first.unwrap();
    assert!(first_result.is_ok(), "first action should parse ok");

    // Second should be hallucination error
    let second = stream.next().await;
    assert!(second.is_some(), "should return hallucination error");
    let err = second.unwrap().unwrap_err();
    assert_eq!(err.code, mad_hatter::InferErrorCode::FormatViolation);
    assert!(err.message.contains("HALLUCINATION"),
        "error should mention HALLUCINATION. Got: {}", err.message);
}

#[tokio::test]
async fn test_stream_planning_at_start_is_not_hallucination() {
    // Planning at the very start (before any action) is normal — not hallucination
    struct NormalPlanningChannel;

    impl LlmChannel for NormalPlanningChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                format!("<planning-{}>\nThinking about this...\n</planning-{}>\n", token, token),
                format!("StreamAction-{}\nreply\ncontent-{}\nDone!\nStreamAction-end-{}\n", token, token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(
        &NormalPlanningChannel,
        &request,
        true,
    ).await.unwrap();

    let first = stream.next().await;
    assert!(first.is_some());
    assert!(first.unwrap().is_ok(), "normal planning at start should not trigger hallucination");

    // Stream should end normally
    let second = stream.next().await;
    assert!(second.is_none(), "stream should end after the single action");
}

#[tokio::test]
async fn test_stream_planning_hallucination_before_first_action() {
    // Planning opens and closes, then another planning tag appears BEFORE any action
    struct DoublePlanningChannel;

    impl LlmChannel for DoublePlanningChannel {
        fn start_stream(&self, prompt: String) -> Result<UnboundedReceiver<String>, String> {
            let token = extract_token(&prompt, "StreamAction");
            Ok(mock_channel_rx(vec![
                // First planning block — normal
                format!("<planning-{}>\nFirst thought\n</planning-{}>\n", token, token),
                // Second planning block — hallucination (planning already completed)
                format!("<planning-{}>\nSecond thought\n</planning-{}>\n", token, token),
            ]))
        }
    }

    let request = StreamReq { prompt: "test".to_string() };
    let mut stream = stream_infer::<StreamReq, StreamAction>(
        &DoublePlanningChannel,
        &request,
        true,
    ).await.unwrap();

    // Should detect hallucination immediately
    let result = stream.next().await;
    assert!(result.is_some(), "should return hallucination error");
    let err = result.unwrap().unwrap_err();
    assert_eq!(err.code, mad_hatter::InferErrorCode::FormatViolation);
    assert!(err.message.contains("HALLUCINATION"),
        "error should mention HALLUCINATION. Got: {}", err.message);
}
