use mem_core::gated_loop::{run_loop, LlmClient, LoopConfig, LoopEvent};
use mem_core::{Chunk, Level, Provenance, Query, Record, Role};
use std::sync::{Arc, Mutex};
use time::OffsetDateTime;

/// Fake LLM that returns scripted responses.
struct FakeLlm {
    responses: Arc<Mutex<Vec<String>>>,
}

impl FakeLlm {
    fn new(responses: Vec<&str>) -> Self {
        Self {
            responses: Arc::new(Mutex::new(responses.iter().map(|s| s.to_string()).collect())),
        }
    }
}

impl LlmClient for FakeLlm {
    fn complete_blocking(&self, _system: &str, _user: &str, _max_tokens: usize) -> anyhow::Result<String> {
        let mut responses = self.responses.lock().unwrap();
        if responses.is_empty() {
            Err(anyhow::Error::msg("No more scripted responses"))
        } else {
            Ok(responses.remove(0))
        }
    }
}

fn make_chunk() -> Chunk {
    Chunk::new(
        1,
        vec![Record {
            role: Role::User,
            text: "test".to_string(),
            timestamp: OffsetDateTime::now_utc(),
            provenance: Provenance {
                source_id: "test".to_string(),
                offset: 0,
            },
        }],
        50,
    )
}

#[test]
fn a1_retain_on_no() {
    let llm = FakeLlm::new(vec![
        "<think>Not useful</think><check>no</check><update>old</update><next>continue</next>";
        5
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk(); 5], &llm).unwrap();
    assert_eq!(outcome.chunks_seen, 5);
    assert_eq!(outcome.chunks_used, 0);
    assert_eq!(outcome.final_memory, "");
}

#[test]
fn a2_update_on_yes() {
    let llm = FakeLlm::new(vec![
        "<think>No</think><check>no</check><update>old</update><next>continue</next>",
        "<think>No</think><check>no</check><update>old</update><next>continue</next>",
        "<think>Yes</think><check>yes</check><update>New memory</update><next>continue</next>",
        "<think>No</think><check>no</check><update>old</update><next>continue</next>",
        "<think>No</think><check>no</check><update>old</update><next>continue</next>",
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk(); 5], &llm).unwrap();
    assert_eq!(outcome.chunks_seen, 5);
    assert_eq!(outcome.chunks_used, 1);
    assert_eq!(outcome.final_memory, "New memory");
}

#[test]
fn a3_exit_gate_off_reads_all() {
    let llm = FakeLlm::new(vec![
        "<think>End</think><check>no</check><update>x</update><next>end</next>";
        10
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk(); 10], &llm).unwrap();
    assert_eq!(outcome.chunks_seen, 10);
}

#[test]
fn a4_exit_gate_on_stops() {
    let llm = FakeLlm::new(vec![
        "<think>End</think><check>no</check><update>x</update><next>end</next>",
        "<think>End</think><check>no</check><update>x</update><next>end</next>",
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: true,
    };

    let outcome = run_loop(config, vec![make_chunk(); 10], &llm).unwrap();
    assert!(outcome.chunks_seen < 10);
}

#[test]
fn a5_exit_always_recorded() {
    let llm = FakeLlm::new(vec![
        "<think>End</think><check>no</check><update>x</update><next>end</next>",
        "<think>End</think><check>no</check><update>x</update><next>end</next>",
        "<think>End</think><check>no</check><update>x</update><next>end</next>",
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk(); 3], &llm).unwrap();
    let exit_count = outcome.events.iter().filter(|e| {
        matches!(e, LoopEvent::Gate { exit: true, .. })
    }).count();
    assert!(exit_count > 0);
}

#[test]
fn a6_budget_exceeded_retains() {
    let large_text = "x".repeat(2000);
    let update_large = format!("<think>Big</think><check>yes</check><update>{}</update><next>continue</next>", large_text);
    
    let llm = FakeLlm::new(vec![&update_large]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
    assert_eq!(outcome.final_memory, "");
    assert_eq!(outcome.chunks_used, 0);
    assert!(outcome.events.iter().any(|e| matches!(e, LoopEvent::BudgetExceeded { .. })));
}

#[test]
fn a7_parse_retry() {
    let llm = FakeLlm::new(vec![
        "malformed",
        "also bad",
        "<think>Good</think><check>yes</check><update>Memory</update><next>continue</next>",
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
    assert_eq!(outcome.final_memory, "Memory");
    assert_eq!(outcome.chunks_used, 1);
}

#[test]
fn a8_parse_failure_continues() {
    let llm = FakeLlm::new(vec![
        "bad",
        "bad",
        "bad",
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
    assert_eq!(outcome.chunks_used, 0);
    assert!(outcome.events.iter().any(|e| matches!(e, LoopEvent::ParseFailed { .. })));
}

#[test]
fn a9_parents_linked() {
    let llm = FakeLlm::new(vec![
        "<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
    ]);

    let config = LoopConfig {
        level: Level::L0,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
    assert!(outcome.events.iter().any(|e| matches!(e, LoopEvent::Evidence { .. })));
}

#[test]
fn a10_level_is_parameter() {
    let llm = FakeLlm::new(vec![
        "<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
        "<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
    ]);

    let chunk = make_chunk();
    
    let config_l1 = LoopConfig {
        level: Level::L1,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let config_l2 = LoopConfig {
        level: Level::L2,
        query: Query {
            id: "test".to_string(),
            question: "Test?".to_string(),
            exit_gate: false,
        },
        memory_budget: 1024,
        use_exit_gate: false,
    };

    let outcome_l1 = run_loop(config_l1, vec![chunk.clone()], &llm).unwrap();
    // Note: LlmClient consumed, so create new for second run
    let llm2 = FakeLlm::new(vec![
        "<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
    ]);
    let outcome_l2 = run_loop(config_l2, vec![chunk.clone()], &llm2).unwrap();

    // Both should have same event count (just different level internally)
    assert_eq!(outcome_l1.events.len(), outcome_l2.events.len());
}