//! Resume-path integration tests: verify infer() handles Truncated by
//! saving state and re-issuing with prefill on the next attempt.
//!
//! Scenario:
//! - Attempt 1: LLM returns a truncated buffer (mid-variant), finish=stop.
//!   StreamInfer detects clean EOF but parse fails → Truncated error.
//! - infer() saves (token, buffer), retries.
//! - Attempt 2: infer() calls start_stream_with_resume with the saved
//!   buffer as prefill. OpenAiChannel embeds it as an assistant-role
//!   message. MockServer returns only the continuation. StreamInfer
//!   reuses saved_buffer, appends new chunks, parse succeeds.
//!
//! These tests lock down the resume contract that production depends on.

mod common;

use common::mock_llm_server::{CapturedRequest, MockLlmServer, Script};
use mad_hatter::{infer, OpenAiChannel, Parseable, Promptable};
use std::sync::Arc;

// ============================================================
// Test types (same shape as happy-path tests)
// ============================================================

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

#[derive(Parseable, PartialEq, Debug)]
enum TestAction {
    #[prompt("Reply to user")]
    Reply { content: String },
    #[prompt("Do nothing")]
    Idle,
}

// ============================================================
// Helpers
// ============================================================

fn extract_action_token(body: &str) -> String {
    let prefix = "TestAction-";
    let mut search_pos = 0;
    while let Some(pos) = body[search_pos..].find(prefix) {
        let abs = search_pos + pos;
        let after = &body[abs + prefix.len()..];
        if after.starts_with("end-") {
            search_pos = abs + prefix.len() + "end-".len();
            continue;
        }
        let tok_end = after
            .find(|c: char| !c.is_ascii_alphanumeric())
            .unwrap_or(after.len());
        if tok_end > 0 {
            return after[..tok_end].to_string();
        }
        search_pos = abs + prefix.len();
    }
    panic!("no TestAction-<token> found in body");
}

/// First-attempt output: only the separator + start of variant name "re".
/// SSE stream closes cleanly (finish=stop) but the buffer cannot parse
/// into TestAction → StreamInfer yields Truncated.
fn build_truncated_output(token: &str) -> String {
    format!("TestAction-{tok}\nre", tok = token)
}

/// Second-attempt continuation: picks up right where truncation stopped.
/// When StreamInfer reuses saved_buffer="TestAction-{tok}\nre" and
/// appends this, the combined buffer is a valid Reply variant.
fn build_resume_continuation(content: &str, token: &str) -> String {
    format!(
        "ply\ncontent-{tok}\n{c}\nend-marker-{tok}\n",
        tok = token,
        c = content
    )
}

// ============================================================
// Tests
// ============================================================

#[tokio::test]
async fn test_infer_resume_after_truncation_recovers() {
    let mock = MockLlmServer::start().await;

    // Script 1: first attempt returns truncated content.
    mock.push_script(Script::Reflective(Arc::new(|req: &CapturedRequest| {
        let token = extract_action_token(&req.body);
        Script::SseContent {
            chunks: vec![build_truncated_output(&token)],
            finish_reason: "stop".to_string(),
        }
    })))
    .await;

    // Script 2: second attempt (with prefill) returns the continuation.
    mock.push_script(Script::Reflective(Arc::new(|req: &CapturedRequest| {
        let token = extract_action_token(&req.body);
        Script::SseContent {
            chunks: vec![build_resume_continuation("recovered", &token)],
            finish_reason: "stop".to_string(),
        }
    })))
    .await;

    let channel = OpenAiChannel::new(mock.chat_url(), "test-model", "test-key");
    let req = TestRequest {
        message: "hi".to_string(),
    };

    let (action, _usage) = infer::<_, TestAction>(&channel, &req, false)
        .await
        .expect("infer should succeed after resume");

    assert_eq!(
        action,
        TestAction::Reply {
            content: "recovered".to_string()
        }
    );

    // Two HTTP requests were made (original + resume).
    assert_eq!(mock.request_count().await, 2);
}

#[tokio::test]
async fn test_infer_resume_sends_prefill_as_assistant_message() {
    let mock = MockLlmServer::start().await;

    mock.push_script(Script::Reflective(Arc::new(|req: &CapturedRequest| {
        let token = extract_action_token(&req.body);
        Script::SseContent {
            chunks: vec![build_truncated_output(&token)],
            finish_reason: "stop".to_string(),
        }
    })))
    .await;

    mock.push_script(Script::Reflective(Arc::new(|req: &CapturedRequest| {
        let token = extract_action_token(&req.body);
        Script::SseContent {
            chunks: vec![build_resume_continuation("ok", &token)],
            finish_reason: "stop".to_string(),
        }
    })))
    .await;

    let channel = OpenAiChannel::new(mock.chat_url(), "test-model", "test-key");
    let req = TestRequest {
        message: "go".to_string(),
    };

    let _ = infer::<_, TestAction>(&channel, &req, false)
        .await
        .expect("infer should succeed");

    let requests = mock.requests().await;
    assert_eq!(requests.len(), 2, "expected 2 HTTP requests");

    // First request: no assistant message prefill.
    let first_body = &requests[0].body;
    assert!(
        !first_body.contains("\"role\":\"assistant\"")
            && !first_body.contains("\"role\": \"assistant\""),
        "first request should not contain assistant role message, got: {}",
        &first_body[..first_body.len().min(500)]
    );

    // Second request: must carry the truncated buffer as assistant prefill.
    let second_body = &requests[1].body;
    assert!(
        second_body.contains("assistant"),
        "second request should contain 'assistant' role, got: {}",
        &second_body[..second_body.len().min(500)]
    );
    // The saved buffer "TestAction-{tok}\nre" should appear verbatim
    // (JSON-encoded, so "\n" becomes "\\n").
    let token = extract_action_token(first_body);
    let buffer_needle = format!("TestAction-{}", token);
    assert!(
        second_body.contains(&buffer_needle),
        "second request should contain saved buffer '{}', got: {}",
        buffer_needle,
        &second_body[..second_body.len().min(500)]
    );
}