pe-graph 0.1.0

//! Integration tests for the graph execution engine.
//!
//! Tests the full BSP cycle: StateGraph → compile → invoke → verify outcome.

use pe_core::error::PeError;
use pe_core::node::{InterruptRequest, NodeContext, NodeFn, NodeFuture, NodeResult};
use pe_core::state::{State, StateUpdate};
use pe_core::types::{END, START};
use pe_graph::{ExecutionOutcome, GraphConfig, InMemoryCheckpointer, StateGraph};
use serde::{Deserialize, Serialize};
use std::time::Duration;

// ── Test Types ────────────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct TestState {
    messages: Vec<String>,
    counter: u32,
    thread_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
struct TestUpdate {
    messages: Option<Vec<String>>,
    counter: Option<u32>,
}

impl StateUpdate for TestUpdate {}

impl State for TestState {
    type Update = TestUpdate;
    fn apply(&mut self, update: TestUpdate) {
        if let Some(msgs) = update.messages {
            self.messages.extend(msgs);
        }
        if let Some(c) = update.counter {
            self.counter = c;
        }
    }
}

impl TestState {
    fn new() -> Self {
        Self {
            messages: Vec::new(),
            counter: 0,
            thread_id: "test".into(),
        }
    }
}

// ── Test Nodes ────────────────────────────────────────────────────────

struct AppendNode {
    node_name: &'static str,
    msg: &'static str,
}

impl AppendNode {
    fn new(name: &'static str, msg: &'static str) -> Self {
        Self {
            node_name: name,
            msg,
        }
    }
}

impl NodeFn<TestState> for AppendNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let msg = self.msg.to_string();
        Box::pin(async move {
            NodeResult::Update(TestUpdate {
                messages: Some(vec![msg]),
                counter: None,
            })
        })
    }
    fn name(&self) -> &str {
        self.node_name
    }
}

struct IncrementNode {
    node_name: &'static str,
}

impl NodeFn<TestState> for IncrementNode {
    fn call(&self, state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let new_val = state.counter + 1;
        Box::pin(async move {
            NodeResult::Update(TestUpdate {
                messages: None,
                counter: Some(new_val),
            })
        })
    }
    fn name(&self) -> &str {
        self.node_name
    }
}

struct InterruptNode;

impl NodeFn<TestState> for InterruptNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        Box::pin(async {
            NodeResult::Interrupt(InterruptRequest {
                reason: "need approval".into(),
                partial_update: None,
                resume_point: "review".into(),
            })
        })
    }
    fn name(&self) -> &str {
        "interrupt"
    }
}

struct ErrorNode;

impl NodeFn<TestState> for ErrorNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        Box::pin(async {
            NodeResult::Error(PeError::Internal {
                details: "node failed".into(),
            })
        })
    }
    fn name(&self) -> &str {
        "error"
    }
}

// ── Tests ─────────────────────────────────────────────────────────────

#[tokio::test]
async fn test_linear_graph_executes_in_order() {
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "first"))
        .add_node("b", AppendNode::new("b", "second"))
        .add_node("c", AppendNode::new("c", "third"))
        .add_edge(START, "a")
        .add_edge("a", "b")
        .add_edge("b", "c")
        .add_edge("c", END)
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages, vec!["first", "second", "third"]);
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_cycle_hits_recursion_limit() {
    // a → b → a (infinite loop)
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "ping"))
        .add_node("b", AppendNode::new("b", "pong"))
        .add_edge(START, "a")
        .add_edge("a", "b")
        .add_edge("b", "a")
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_recursion_limit(5);
    let result = graph.invoke(TestState::new(), config).await;

    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        matches!(err, PeError::GraphRecursion { limit: 5 }),
        "Expected GraphRecursion, got: {err:?}"
    );
}

#[tokio::test]
async fn test_conditional_edge_routes_correctly() {
    // chat → tools (if counter < 2) or END (if counter >= 2)
    let graph = StateGraph::new()
        .add_node("chat", IncrementNode { node_name: "chat" })
        .add_node("tools", AppendNode::new("tools", "tool_call"))
        .add_edge(START, "chat")
        .add_conditional_edge("chat", |state: &TestState| {
            if state.counter < 2 {
                vec!["tools".into()]
            } else {
                vec![END.into()]
            }
        })
        .add_edge("tools", "chat")
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // chat runs 3 times (counter: 0→1, 1→2, at 2 → END)
            // tools runs 2 times (when counter was 1 and 2... wait let me trace:
            // Step 1: chat runs (counter 0→1), conditional: 1 < 2 → tools
            // Step 2: tools runs (appends "tool_call"), then edge tools→chat
            // Step 3: chat runs (counter 1→2), conditional: 2 >= 2 → END
            assert_eq!(state.counter, 2);
            assert_eq!(state.messages, vec!["tool_call"]);
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_parallel_edges_both_execute() {
    // START → a, START → b (both run in same superstep)
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "from_a"))
        .add_node("b", AppendNode::new("b", "from_b"))
        .add_edge(START, "a")
        .add_edge(START, "b")
        .add_edge("a", END)
        .add_edge("b", END)
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // Both messages should be present (Appender semantics)
            assert_eq!(state.messages.len(), 2);
            assert!(state.messages.contains(&"from_a".to_string()));
            assert!(state.messages.contains(&"from_b".to_string()));
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_interrupt_halts_execution() {
    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "done"))
        .add_node("review", InterruptNode)
        .add_edge(START, "work")
        .add_edge("work", "review")
        .add_edge("review", END)
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Interrupted { state, request } => {
            // "work" completed before interrupt
            assert_eq!(state.messages, vec!["done"]);
            assert_eq!(request.reason, "need approval");
            assert_eq!(request.resume_point, "review");
        }
        _ => panic!("Expected Interrupted"),
    }
}

#[tokio::test]
async fn test_error_propagates() {
    let graph = StateGraph::new()
        .add_node("bad", ErrorNode)
        .add_edge(START, "bad")
        .add_edge("bad", END)
        .compile()
        .unwrap();

    let result = graph.invoke(TestState::new(), GraphConfig::default()).await;
    assert!(result.is_err());
    let err = result.unwrap_err().to_string();
    assert!(err.contains("node failed"), "Error should propagate: {err}");
}

#[tokio::test]
async fn test_snapshot_isolation() {
    // Two parallel nodes both read counter=0 and try to set counter=1
    // With snapshot isolation, both see counter=0
    // LastValue semantics means last-applied wins (both set to 1, not 2)
    struct SetCounterNode {
        name: &'static str,
    }

    impl NodeFn<TestState> for SetCounterNode {
        fn call(&self, state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
            // Read current counter and add 1
            let val = state.counter + 1;
            Box::pin(async move {
                NodeResult::Update(TestUpdate {
                    messages: None,
                    counter: Some(val),
                })
            })
        }
        fn name(&self) -> &str {
            self.name
        }
    }

    let graph = StateGraph::new()
        .add_node("a", SetCounterNode { name: "a" })
        .add_node("b", SetCounterNode { name: "b" })
        .add_edge(START, "a")
        .add_edge(START, "b")
        .add_edge("a", END)
        .add_edge("b", END)
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // Both saw counter=0, both set counter=1
            // LastValue: last applied wins = 1 (not 2)
            assert_eq!(
                state.counter, 1,
                "Snapshot isolation: both nodes saw counter=0"
            );
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_snapshot_isolation_appender() {
    // Two parallel nodes both append messages — Appender semantics means BOTH preserved
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "from_a"))
        .add_node("b", AppendNode::new("b", "from_b"))
        .add_edge(START, "a")
        .add_edge(START, "b")
        .add_edge("a", END)
        .add_edge("b", END)
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // Appender: both writes accumulate (unlike LastValue which overwrites)
            assert_eq!(state.messages.len(), 2);
            assert!(state.messages.contains(&"from_a".to_string()));
            assert!(state.messages.contains(&"from_b".to_string()));
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_node_context_remaining_steps() {
    // Node reads ctx.remaining_steps() and puts it in the message
    struct ContextAwareNode;

    impl NodeFn<TestState> for ContextAwareNode {
        fn call(&self, _state: &TestState, ctx: &NodeContext) -> NodeFuture<TestUpdate> {
            let remaining = ctx.remaining_steps();
            let is_last = ctx.is_last_step();
            let msg = format!("remaining={remaining},last={is_last}");
            Box::pin(async move {
                NodeResult::Update(TestUpdate {
                    messages: Some(vec![msg]),
                    counter: None,
                })
            })
        }
        fn name(&self) -> &str {
            "ctx-aware"
        }
    }

    let graph = StateGraph::new()
        .add_node("check", ContextAwareNode)
        .add_edge(START, "check")
        .add_edge("check", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_recursion_limit(10);
    let outcome = graph.invoke(TestState::new(), config).await.unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // Step 1, limit 10 → remaining = 10 - 1 = 9, not last
            assert_eq!(state.messages, vec!["remaining=9,last=false"]);
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_get_state_history() {
    let cp = InMemoryCheckpointer::new();

    // Graph with 3 linear steps → 3 checkpoints
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "first"))
        .add_node("b", AppendNode::new("b", "second"))
        .add_node("c", AppendNode::new("c", "third"))
        .add_edge(START, "a")
        .add_edge("a", "b")
        .add_edge("b", "c")
        .add_edge("c", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

    let config = GraphConfig::default().with_thread_id("history-test");
    let _ = graph.invoke(TestState::new(), config).await.unwrap();

    let history = graph.get_state_history("history-test").await.unwrap();
    // Should have checkpoints (one per superstep)
    assert!(
        history.len() >= 3,
        "Expected at least 3 checkpoints, got {}",
        history.len()
    );
    // Steps should be monotonically increasing
    for i in 1..history.len() {
        assert!(history[i].step >= history[i - 1].step);
    }
}

#[tokio::test]
async fn test_resume_after_interrupt() {
    let cp = InMemoryCheckpointer::new();

    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_node("review", InterruptNode)
        .add_node("finish", AppendNode::new("finish", "step2"))
        .add_edge(START, "work")
        .add_edge("work", "review")
        .add_edge("review", "finish")
        .add_edge("finish", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

    let config = GraphConfig::default().with_thread_id("resume-test");

    // First invoke — hits interrupt
    let outcome = graph
        .invoke(TestState::new(), config.clone())
        .await
        .unwrap();
    assert!(matches!(outcome, ExecutionOutcome::Interrupted { .. }));

    // Resume with human input
    let human_input = TestUpdate {
        messages: Some(vec!["human says ok".into()]),
        counter: None,
    };
    let outcome = graph
        .resume("resume-test", human_input, config)
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"step1".to_string()));
            assert!(state.messages.contains(&"human says ok".to_string()));
            // "finish" node should have run after resume
            assert!(state.messages.contains(&"step2".to_string()));
        }
        _ => panic!("Expected Completed after resume"),
    }
}

#[tokio::test]
async fn test_get_state_after_checkpoint() {
    let cp = InMemoryCheckpointer::new();

    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "data"))
        .add_node("review", InterruptNode)
        .add_edge(START, "work")
        .add_edge("work", "review")
        .add_edge("review", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

    let config = GraphConfig::default().with_thread_id("state-test");
    let _ = graph.invoke(TestState::new(), config).await.unwrap();

    let snapshot = graph.get_state("state-test").await.unwrap();
    assert!(snapshot.is_some());

    let snap = snapshot.unwrap();
    assert!(snap.state.messages.contains(&"data".to_string()));
    assert_eq!(snap.thread_id, "state-test");
}

#[tokio::test]
async fn test_empty_graph_no_start_edge() {
    let result = StateGraph::<TestState>::new()
        .add_node("orphan", AppendNode::new("orphan", "lost"))
        .compile();

    assert!(result.is_err());
}

#[tokio::test]
async fn test_single_node_graph() {
    let graph = StateGraph::new()
        .add_node("only", AppendNode::new("only", "hello"))
        .add_edge(START, "only")
        .add_edge("only", END)
        .compile()
        .unwrap();

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages, vec!["hello"]);
        }
        _ => panic!("Expected Completed"),
    }
}

// ── Timeout Tests ────────────────────────────────────────────────────

/// A node that sleeps for a specified duration before returning.
struct SlowNode {
    name: &'static str,
    delay: Duration,
}

impl NodeFn<TestState> for SlowNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let delay = self.delay;
        Box::pin(async move {
            tokio::time::sleep(delay).await;
            NodeResult::Update(TestUpdate {
                messages: Some(vec!["slow_done".into()]),
                counter: None,
            })
        })
    }
    fn name(&self) -> &str {
        self.name
    }
}

#[tokio::test]
async fn test_no_timeout_runs_normally() {
    // Default config has no timeout — graph completes without issue
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "first"))
        .add_node("b", AppendNode::new("b", "second"))
        .add_edge(START, "a")
        .add_edge("a", "b")
        .add_edge("b", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default();
    assert!(config.max_execution_time.is_none());

    let outcome = graph.invoke(TestState::new(), config).await.unwrap();
    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages, vec!["first", "second"]);
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_generous_timeout_completes_normally() {
    // 10-second timeout on a fast graph — should complete with time to spare
    let graph = StateGraph::new()
        .add_node("a", AppendNode::new("a", "done"))
        .add_edge(START, "a")
        .add_edge("a", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_max_execution_time(Duration::from_secs(10));

    let outcome = graph.invoke(TestState::new(), config).await.unwrap();
    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages, vec!["done"]);
        }
        _ => panic!("Expected Completed"),
    }
}

#[tokio::test]
async fn test_short_timeout_on_multi_step_graph_returns_timeout() {
    // A cycling graph with a 1ms timeout — the first superstep may complete
    // (the timeout check happens at the START of each superstep), but subsequent
    // iterations will exceed the 1ms budget and return Timeout.
    let graph = StateGraph::new()
        .add_node(
            "work",
            SlowNode {
                name: "work",
                delay: Duration::from_millis(10),
            },
        )
        .add_node("loop_back", AppendNode::new("loop_back", "again"))
        .add_edge(START, "work")
        .add_edge("work", "loop_back")
        .add_edge("loop_back", "work") // cycle
        .compile()
        .unwrap();

    let config = GraphConfig::default()
        .with_recursion_limit(100) // high limit so timeout triggers first
        .with_max_execution_time(Duration::from_millis(1));

    let result = graph.invoke(TestState::new(), config).await;
    assert!(result.is_err(), "Expected timeout error");
    let err = result.unwrap_err();
    assert!(
        matches!(err, PeError::Timeout { .. }),
        "Expected PeError::Timeout, got: {err:?}"
    );
}

#[tokio::test]
async fn test_timeout_error_contains_elapsed_info() {
    // Verify the Timeout variant has meaningful data (seconds field)
    let graph = StateGraph::new()
        .add_node(
            "slow",
            SlowNode {
                name: "slow",
                delay: Duration::from_millis(50),
            },
        )
        .add_node("next", AppendNode::new("next", "x"))
        .add_edge(START, "slow")
        .add_edge("slow", "next")
        .add_edge("next", "slow") // cycle
        .compile()
        .unwrap();

    let config = GraphConfig::default()
        .with_recursion_limit(100)
        .with_max_execution_time(Duration::from_millis(1));

    let result = graph.invoke(TestState::new(), config).await;
    let err = result.unwrap_err();
    match err {
        PeError::Timeout { seconds } => {
            // The Display message should contain "timeout"
            let msg = format!("{err}");
            assert!(
                msg.to_lowercase().contains("timeout"),
                "Error message should mention timeout: {msg}"
            );
            // seconds should be small since we're timing out after ~50ms
            assert!(
                seconds < 5.0,
                "Sub-second timeout should report small value, got {seconds}"
            );
        }
        other => panic!("Expected PeError::Timeout, got: {other:?}"),
    }
}