use mem_core::prompt::PromptBuilder;
use mem_core::{Chunk, PromptMessages, Query, Record, Role, Provenance};
use time::OffsetDateTime;

fn make_chunk(records: Vec<(Role, &str)>) -> Chunk {
    let records = records
        .into_iter()
        .map(|(role, text)| Record {
            role,
            text: text.to_string(),
            timestamp: OffsetDateTime::now_utc(),
            provenance: Provenance {
                source_id: "test".to_string(),
                offset: 0,
            },
        })
        .collect();
    
    Chunk::new(1, records, 100)
}

#[test]
fn a1_golden_t1() {
    let query = Query {
        id: "architecture-decisions".to_string(),
        question: "What architectural decisions were made?".to_string(),
        exit_gate: false,
    };

    let chunk = make_chunk(vec![
        (Role::User, "Tell me about the architecture"),
        (Role::Assistant, "We use a microservices design"),
    ]);

    let (_system, user) = PromptBuilder::build(&query, None, &chunk).expect("Should build");

    // Load golden file
    let golden = std::fs::read_to_string("fixtures/expected/prompt-t1.txt")
        .expect("Should read golden file");

    assert_eq!(user, golden, "User message should match golden file exactly");
}

#[test]
fn a2_golden_tn() {
    let query = Query {
        id: "architecture-decisions".to_string(),
        question: "What architectural decisions were made?".to_string(),
        exit_gate: false,
    };

    let chunk = make_chunk(vec![
        (Role::User, "What about the database?"),
        (Role::Assistant, "We chose PostgreSQL for primary storage."),
    ]);

    let prior_memory = "We use a microservices design with REST APIs.";
    let (_system, user) = PromptBuilder::build(&query, Some(prior_memory), &chunk)
        .expect("Should build");

    let golden = std::fs::read_to_string("fixtures/expected/prompt-tn.txt")
        .expect("Should read golden file");

    assert_eq!(user, golden, "User message should match golden file exactly");
}

#[test]
fn a3_no_previous_memory_literal() {
    let query = Query {
        id: "test".to_string(),
        question: "Test question?".to_string(),
        exit_gate: false,
    };

    let chunk = make_chunk(vec![(Role::User, "Small chunk")]);

    let (_system, user) = PromptBuilder::build(&query, None, &chunk).expect("Should build");

    // At t=1, should contain the literal string "No previous memory"
    assert!(
        user.contains("No previous memory"),
        "t=1 prompt should contain 'No previous memory' literally"
    );
}

#[test]
fn a4_all_tags_present() {
    let query = Query {
        id: "test".to_string(),
        question: "Test question?".to_string(),
        exit_gate: false,
    };

    let chunk = make_chunk(vec![(Role::User, "chunk content")]);

    let (_system, user) = PromptBuilder::build(&query, None, &chunk).expect("Should build");

    // All three tags should appear exactly once
    assert_eq!(
        user.matches("<problem>").count(),
        1,
        "<problem> should appear exactly once"
    );
    assert_eq!(
        user.matches("</problem>").count(),
        1,
        "</problem> should appear exactly once"
    );
    assert_eq!(
        user.matches("<memory>").count(),
        1,
        "<memory> should appear exactly once"
    );
    assert_eq!(
        user.matches("</memory>").count(),
        1,
        "</memory> should appear exactly once"
    );
    assert_eq!(
        user.matches("<section>").count(),
        1,
        "<section> should appear exactly once"
    );
    assert_eq!(
        user.matches("</section>").count(),
        1,
        "</section> should appear exactly once"
    );
}

#[test]
fn a5_role_labels_rendered() {
    let query = Query {
        id: "test".to_string(),
        question: "Test?".to_string(),
        exit_gate: false,
    };

    let chunk = make_chunk(vec![
        (Role::User, "User says something"),
        (Role::Assistant, "Assistant responds"),
        (Role::ToolResult, "Tool feedback"),
        (Role::System, "System message"),
    ]);

    let (_system, user) = PromptBuilder::build(&query, None, &chunk).expect("Should build");

    assert!(user.contains("[User]"), "Should contain [User] label");
    assert!(
        user.contains("[Assistant]"),
        "Should contain [Assistant] label"
    );
    assert!(
        user.contains("[ToolResult]"),
        "Should contain [ToolResult] label"
    );
    assert!(user.contains("[System]"), "Should contain [System] label");
}

#[test]
fn a6_over_budget_errors() {
    let query = Query {
        id: "test".to_string(),
        question: "Test question?".to_string(),
        exit_gate: false,
    };

    // Create a very large chunk that exceeds budget
    let huge_content = "x".repeat(6000); // Over 5000 budget
    let chunk = make_chunk(vec![(Role::User, &huge_content)]);

    let result = PromptBuilder::build(&query, None, &chunk);
    assert!(result.is_err(), "Should error on over-budget chunk");

    let err_msg = format!("{:?}", result.err().unwrap());
    assert!(
        err_msg.contains("Chunk budget") || err_msg.contains("section"),
        "Error should mention chunk/section budget"
    );
}

#[test]
fn a7_budget_headroom() {
    let query = Query {
        id: "test".to_string(),
        question: "Test question?".to_string(),
        exit_gate: false,
    };

    // Create a realistic chunk (under budget)
    let chunk_content = "x".repeat(4000); // Under 5000 budget
    let chunk = make_chunk(vec![(Role::User, &chunk_content)]);

    let (system, user) = PromptBuilder::build(&query, None, &chunk)
        .expect("Should build under-budget prompt");

    // Rough estimate: 4 chars ≈ 1 token
    let total_size = system.len() + user.len();
    let tokens_estimate = total_size / 4;

    // Should have headroom: 32768 - 2048 (response) = 30720 available
    assert!(
        tokens_estimate < 30720 - 100, // 100 token safety margin
        "Should have headroom for response: {} tokens used, {} available",
        tokens_estimate,
        30720
    );
}

// ── Cache alignment integration tests ──────────────────────────────────

#[test]
fn a8_cache_prefix_stable_across_50_chunks() {
    // Simulate a real ingestion run: same query, 50 different chunks.
    // The cache prefix (system + query) must be identical every time.
    let query = Query {
        id: "arch-decisions".to_string(),
        question: "What architectural decisions were made and why?".to_string(),
        exit_gate: false,
    };

    let mut prefixes: Vec<(String, String)> = Vec::new();

    for i in 0..50 {
        let content = format!("Evidence chunk number {} with unique content {}", i, "x".repeat(i * 10));
        let chunk = make_chunk(vec![(Role::User, &content)]);
        let memory = if i == 0 {
            None
        } else {
            Some("Accumulated memory from previous turns")
        };

        let msgs = PromptBuilder::build_cache_aligned(&query, memory, &chunk)
            .expect("Should build cache-aligned prompt");

        prefixes.push((msgs.system.clone(), msgs.user_messages[0].clone()));
    }

    // All 50 system messages must be identical
    for (i, (system, _)) in prefixes.iter().enumerate() {
        assert_eq!(
            system, &prefixes[0].0,
            "System message must be stable (chunk {})", i
        );
    }

    // All 50 query messages must be identical
    for (i, (_, query_msg)) in prefixes.iter().enumerate() {
        assert_eq!(
            query_msg, &prefixes[0].1,
            "Query message must be stable (chunk {})", i
        );
    }
}

#[test]
fn a9_cache_aligned_headroom() {
    // Cache-aligned prompts should also have positive headroom
    let query = Query {
        id: "test".to_string(),
        question: "Test question?".to_string(),
        exit_gate: false,
    };

    let chunk_content = "x".repeat(4000);
    let chunk = make_chunk(vec![(Role::User, &chunk_content)]);

    let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk)
        .expect("Should build cache-aligned prompt");

    assert!(
        msgs.headroom() > 0,
        "Cache-aligned prompt should have positive headroom: {} tokens",
        msgs.headroom()
    );

    // Cache prefix should be meaningful (not zero)
    assert!(
        msgs.cache_prefix_tokens() > 50,
        "Cache prefix should be >50 tokens, got {}",
        msgs.cache_prefix_tokens()
    );
}

#[test]
fn a10_cache_savings_estimate() {
    // Demonstrate the savings: compare total tokens vs cached tokens
    // across a simulated 20-chunk run
    let query = Query {
        id: "test".to_string(),
        question: "What are the deployment patterns?".to_string(),
        exit_gate: false,
    };

    let mut total_input_tokens = 0usize;
    let mut total_cached_tokens = 0usize;

    for i in 0..20 {
        let content = format!("Deployment evidence chunk {} with details", i);
        let chunk = make_chunk(vec![(Role::User, &content)]);

        let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk)
            .expect("Should build");

        total_input_tokens += msgs.total_tokens();
        total_cached_tokens += msgs.cache_prefix_tokens();
    }

    // Cache prefix should be a significant portion of total input
    let cache_ratio = total_cached_tokens as f64 / total_input_tokens as f64;
    assert!(
        cache_ratio > 0.3,
        "Cache ratio should be >30%, got {:.1}% ({}/{} tokens)",
        cache_ratio * 100.0,
        total_cached_tokens,
        total_input_tokens
    );
}