//! Disconnect-path integration tests: verify infer() handles mid-stream
//! network drops (TCP close without [DONE]/finish_reason) by classifying
//! as Truncated and attempting resume.
//!
//! Contract covered:
//! - SSE stream closed mid-transmission → `InferErrorCode::Truncated`
//! - `Truncated` is retryable (`is_retryable() == true`)
//! - After `INFER_MAX_RETRIES + 1 = 3` failed attempts, final error surfaces
//! - If a later attempt receives a complete response, infer succeeds via resume

mod common;

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

// ============================================================
// Test types
// ============================================================

#[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");
}

fn build_complete_idle(token: &str) -> String {
    format!(
        "TestAction-{tok}\nidle\nend-marker-{tok}\n",
        tok = token
    )
}

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

/// 3 consecutive abrupt disconnects (no [DONE], no finish_reason) exhaust
/// the retry budget and surface the Truncated error.
#[tokio::test]
async fn test_infer_disconnect_exhausts_retries_as_truncated() {
    let mock = MockLlmServer::start().await;

    // Push 3 scripts — each ends the stream abruptly mid-transmission.
    for _ in 0..3 {
        mock.push_script(Script::SseAbruptClose {
            chunks: vec!["TestAction-".to_string()],
        })
        .await;
    }

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

    let err = infer::<_, TestAction>(&channel, &req, false)
        .await
        .expect_err("infer should exhaust retries");

    assert_eq!(
        err.code,
        InferErrorCode::Truncated,
        "final error should be Truncated, got: {:?}",
        err
    );

    // Confirm all 3 attempts were made.
    assert_eq!(
        mock.request_count().await,
        3,
        "expected 3 HTTP requests (1 original + 2 retries)"
    );
}

/// First attempt drops mid-stream; second attempt completes cleanly.
/// infer() should succeed via the Truncated → resume retry path.
#[tokio::test]
async fn test_infer_disconnect_recovers_on_retry() {
    let mock = MockLlmServer::start().await;

    // Attempt 1: abrupt close mid-stream (empty buffer, truncated).
    mock.push_script(Script::SseAbruptClose { chunks: vec![] })
        .await;

    // Attempt 2: complete idle response.
    mock.push_script(Script::Reflective(Arc::new(|req: &CapturedRequest| {
        let token = extract_action_token(&req.body);
        Script::SseContent {
            chunks: vec![build_complete_idle(&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 (action, _usage) = infer::<_, TestAction>(&channel, &req, false)
        .await
        .expect("infer should recover after disconnect");

    assert_eq!(action, TestAction::Idle);
    assert_eq!(mock.request_count().await, 2);
}

/// Truncated is correctly classified as retryable — locks the retry policy contract.
#[tokio::test]
async fn test_truncated_error_is_retryable() {
    assert!(InferErrorCode::Truncated.is_retryable());
    assert!(!InferErrorCode::Truncated.is_channel());
    assert!(!InferErrorCode::Truncated.is_quota());
}