use mem_core::prompt::PromptBuilder;
use mem_core::{Chunk, 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
    );
}