//! Test: infer resume (breakpoint continuation) on truncation.
//!
//! Verifies that when a stream is truncated mid-output, the retry
//! reuses the partial buffer as assistant prefill so the LLM can
//! continue from where it left off.

use mad_hatter::{CachedPrompt, InferError, LlmChannel};
use mad_hatter::llm::{Reflectable, Promptable, Parseable};
use std::sync::atomic::{AtomicUsize, Ordering};

// ---------------------------------------------------------------------------
// Minimal Promptable / Parseable types for testing
// ---------------------------------------------------------------------------

/// Marker trait (no methods).
struct MockReq;
impl Reflectable for MockReq {}

impl Promptable for MockReq {
    fn to_markdown_depth(&self, _depth: usize) -> String {
        "test request".to_string()
    }
}

struct MockResp {
    value: String,
}

impl Reflectable for MockResp {}

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

    fn schema_markdown(token: &str) -> String
    where
        Self: Sized,
    {
        format!(
            "Output exactly in this format:\nMockResp-{token}\nvalue: <your answer>\nMockResp-end-{token}"
        )
    }
}

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

    fn from_markdown(text: &str, token: &str) -> Result<Vec<Self>, String>
    where
        Self: Sized,
    {
        let sep = format!("MockResp-{}", token);
        let end = format!("MockResp-end-{}", token);
        let sep_pos = text.find(&sep).ok_or_else(|| "separator not found".to_string())?;
        let after_sep = sep_pos + sep.len();
        let end_pos = text.find(&end).ok_or_else(|| "end marker not found".to_string())?;
        let content = text[after_sep..end_pos].trim();
        let val = content
            .strip_prefix("value: ")
            .ok_or_else(|| format!("expected 'value: ' prefix, got: {}", content))?;
        Ok(vec![MockResp {
            value: val.to_string(),
        }])
    }
}

// ---------------------------------------------------------------------------
// Mock LLM channel that simulates truncation + resume
// ---------------------------------------------------------------------------

struct ResumeMockChannel {
    /// How many times start_stream_cached was called.
    cached_calls: AtomicUsize,
    /// How many times start_stream_with_resume was called.
    resume_calls: AtomicUsize,
}

impl ResumeMockChannel {
    fn new() -> Self {
        Self {
            cached_calls: AtomicUsize::new(0),
            resume_calls: AtomicUsize::new(0),
        }
    }

    /// Extract the 6-char hex token from a text that contains "MockResp-XXXXXX".
    fn extract_token(text: &str) -> Option<String> {
        let needle = "MockResp-";
        if let Some(pos) = text.find(needle) {
            let after = &text[pos + needle.len()..];
            // Token is 6 hex chars; stop at newline or non-hex
            let candidate: String = after.chars().take(6).collect();
            if candidate.len() == 6 && candidate.chars().all(|c| c.is_ascii_hexdigit()) {
                return Some(candidate);
            }
        }
        None
    }
}

impl LlmChannel for ResumeMockChannel {
    fn start_stream(
        &self,
        _prompt: String,
    ) -> Result<tokio::sync::mpsc::UnboundedReceiver<String>, String> {
        // Should not be called directly; start_stream_cached is used.
        unreachable!("start_stream should not be called directly")
    }

    fn start_stream_cached(
        &self,
        parts: CachedPrompt,
    ) -> Result<tokio::sync::mpsc::UnboundedReceiver<String>, String> {
        let n = self.cached_calls.fetch_add(1, Ordering::SeqCst);
        let merged = parts.merge();
        let token = Self::extract_token(&merged)
            .unwrap_or_else(|| panic!("Cannot extract token from prompt: {}", &merged[..80.min(merged.len())]));

        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

        if n == 0 {
            // First call: send partial output then close channel (truncation).
            tokio::spawn(async move {
                // separator line
                tx.send(format!("MockResp-{}\n", token)).ok();
                // partial content — missing "lo world" and end marker
                tx.send("value: hel".to_string()).ok();
                // drop tx → receiver sees channel closed → Truncated
            });
        } else {
            // Unexpected second call to start_stream_cached (not resume).
            // Send full output as fallback.
            tokio::spawn(async move {
                tx.send(format!("MockResp-{}\n", token)).ok();
                tx.send("value: hello world\n".to_string()).ok();
                tx.send(format!("MockResp-end-{}\n", token)).ok();
            });
        }

        Ok(rx)
    }

    fn start_stream_with_resume(
        &self,
        _parts: CachedPrompt,
        prefill: &str,
    ) -> Result<tokio::sync::mpsc::UnboundedReceiver<String>, String> {
        self.resume_calls.fetch_add(1, Ordering::SeqCst);

        let token = Self::extract_token(prefill)
            .unwrap_or_else(|| panic!("Cannot extract token from prefill: {}", &prefill[..80.min(prefill.len())]));

        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

        // Send the continuation: complete "hel" → "hello world" + end marker
        tokio::spawn(async move {
            tx.send(format!("lo world\nMockResp-end-{}\n", token)).ok();
        });

        Ok(rx)
    }
}

// ---------------------------------------------------------------------------
// Test
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_infer_resume_on_truncation() {
    let channel = ResumeMockChannel::new();

    let result: Result<(MockResp, _), InferError> =
        mad_hatter::infer::<MockReq, MockResp>(&channel, &MockReq).await;

    match result {
        Ok((resp, _usage)) => {
            assert_eq!(resp.value, "hello world", "Resume should produce complete output");
        }
        Err(e) => {
            panic!("Expected successful resume, got error: {} (code: {:?})", e, e.code);
        }
    }

    // start_stream_cached called exactly once (first attempt)
    assert_eq!(
        channel.cached_calls.load(Ordering::SeqCst),
        1,
        "start_stream_cached should be called once (initial attempt)"
    );

    // start_stream_with_resume called exactly once (resume attempt)
    assert_eq!(
        channel.resume_calls.load(Ordering::SeqCst),
        1,
        "start_stream_with_resume should be called once (resume attempt)"
    );
}