pe-graph 0.1.0

//! Integration tests for the interrupt system (Plan 010).
//!
//! Tests the new Command-based resume, PhaseStateStore checkpoint round-trip,
//! and RetryPolicy integration.

use pe_core::node::{HumanInput, InterruptRequest, NodeContext, NodeFn, NodeFuture, NodeResult};
use pe_core::phase_store::PhaseStateStore;
use pe_core::state::{State, StateUpdate};
use pe_core::types::{END, START};
use pe_graph::command::Command;
use pe_graph::retry::{RetryPolicy, with_retry};
use pe_graph::{ExecutionOutcome, GraphConfig, InMemoryCheckpointer, StateGraph};
use serde::{Deserialize, Serialize};

// -- 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,
    message: &'static str,
}

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

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

/// Node that interrupts on first call, continues on resume (when HumanInput is available).
struct ReviewNode;

impl NodeFn<TestState> for ReviewNode {
    fn call(&self, _state: &TestState, ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let has_input = ctx.phase_store.get::<HumanInput>().ok().flatten();
        Box::pin(async move {
            match has_input {
                Some(input) => {
                    // Resumed with human input -- continue
                    let msg = if input.approved {
                        "review-approved"
                    } else {
                        "review-rejected"
                    };
                    NodeResult::Update(TestUpdate {
                        messages: Some(vec![msg.to_string()]),
                        counter: None,
                    })
                }
                None => {
                    // First run -- interrupt for review
                    NodeResult::Interrupt(InterruptRequest {
                        reason: "Review needed".into(),
                        partial_update: Some(TestUpdate {
                            messages: Some(vec!["partial-before-interrupt".into()]),
                            counter: None,
                        }),
                        resume_point: "review:0".into(),
                    })
                }
            }
        })
    }
    fn name(&self) -> &str {
        "review"
    }
}

/// Node that reads phase from PhaseStateStore and behaves differently per phase.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
enum WorkPhase {
    Gathering,
    Processing { data: String },
}

struct PhaseAwareNode;

impl NodeFn<TestState> for PhaseAwareNode {
    fn call(&self, _state: &TestState, ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let phase: Option<WorkPhase> = ctx.phase_store.get::<WorkPhase>().ok().flatten();
        Box::pin(async move {
            match phase {
                None | Some(WorkPhase::Gathering) => {
                    // First run: interrupt to get human input
                    NodeResult::Interrupt(InterruptRequest {
                        reason: "Need data to process".into(),
                        partial_update: Some(TestUpdate {
                            messages: Some(vec!["gathered".into()]),
                            counter: None,
                        }),
                        resume_point: "phase_aware:0".into(),
                    })
                }
                Some(WorkPhase::Processing { data }) => {
                    // Resumed: use the data
                    NodeResult::Update(TestUpdate {
                        messages: Some(vec![format!("processed:{}", data)]),
                        counter: None,
                    })
                }
            }
        })
    }
    fn name(&self) -> &str {
        "phase_aware"
    }
}

// -- Tests --------------------------------------------------------------------

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

    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_node("review", ReviewNode)
        .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("cmd-resume");

    // Invoke -- hits interrupt at review node
    let outcome = graph
        .invoke(TestState::new(), config.clone())
        .await
        .unwrap();

    match &outcome {
        ExecutionOutcome::Interrupted { request, .. } => {
            assert_eq!(request.reason, "Review needed");
            assert_eq!(request.resume_point, "review:0");
        }
        _ => panic!("Expected Interrupted"),
    }

    // Resume with Command::Resume
    let cmd = Command::resume(HumanInput {
        approved: true,
        feedback: Some("approved".into()),
        data: None,
    });

    let outcome = graph.resume_with("cmd-resume", cmd, config).await.unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // step1 ran, partial update applied, review re-ran with input, finish ran
            assert!(state.messages.contains(&"step1".to_string()));
            assert!(
                state
                    .messages
                    .contains(&"partial-before-interrupt".to_string())
            );
            assert!(state.messages.contains(&"review-approved".to_string()));
            assert!(state.messages.contains(&"step2".to_string()));
        }
        _ => panic!("Expected Completed after resume"),
    }
}

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

    // "alt" is reachable via conditional edge from "review" so graph validates.
    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_node("review", ReviewNode)
        .add_node("finish", AppendNode::new("finish", "step2"))
        .add_node("alt", AppendNode::new("alt", "alt-path"))
        .add_edge(START, "work")
        .add_edge("work", "review")
        .add_edge("review", "finish")
        .add_edge("finish", END)
        .add_conditional_edge("review", |_state: &TestState| vec!["alt".into()])
        .add_edge("alt", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

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

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

    // Resume with Command::Goto -- jump directly to 'alt' node
    let cmd = Command::goto("alt");
    let outcome = graph.resume_with("goto-test", cmd, config).await.unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"step1".to_string()));
            assert!(state.messages.contains(&"alt-path".to_string()));
            // 'finish' should NOT have run (goto bypassed normal flow)
            assert!(!state.messages.contains(&"step2".to_string()));
        }
        _ => panic!("Expected Completed after goto"),
    }
}

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

    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_node("review", ReviewNode)
        .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("update-test");

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

    // Resume with Command::Update -- inject state directly
    let cmd = Command::update(serde_json::json!({
        "messages": ["injected-by-command"],
        "counter": 42
    }));
    let outcome = graph.resume_with("update-test", cmd, config).await.unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"injected-by-command".to_string()));
            assert!(state.messages.contains(&"step2".to_string()));
            assert_eq!(state.counter, 42);
        }
        _ => panic!("Expected Completed after update"),
    }
}

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

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

    let config = GraphConfig::default().with_thread_id("bad-goto");

    // Invoke -- hits interrupt
    graph
        .invoke(TestState::new(), config.clone())
        .await
        .unwrap();

    // Goto nonexistent node — must return PeError::GraphValue
    let cmd = Command::goto("nonexistent");
    let result = graph.resume_with("bad-goto", cmd, config).await;
    assert!(result.is_err());
    let err = result.unwrap_err();
    match err {
        pe_core::error::PeError::GraphValue { details } => {
            assert!(
                details.contains("nonexistent"),
                "Error should name the missing node: {details}"
            );
        }
        other => panic!("Expected PeError::GraphValue, got: {other:?}"),
    }
}

#[tokio::test]
async fn test_phase_state_store_survives_checkpoint() {
    // Verify PhaseStateStore serializes through bincode round-trip
    let mut store = PhaseStateStore::new();

    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
    enum MyPhase {
        Init,
        Working { data: String },
    }

    store
        .set(&MyPhase::Working {
            data: "important".into(),
        })
        .unwrap();

    // Bincode round-trip (same as checkpoint serialization)
    let bytes = bincode::serialize(&store).unwrap();
    let restored: PhaseStateStore = bincode::deserialize(&bytes).unwrap();

    let phase: MyPhase = restored.get::<MyPhase>().unwrap().unwrap();
    assert_eq!(
        phase,
        MyPhase::Working {
            data: "important".into()
        }
    );
}

#[tokio::test]
async fn test_retry_within_node() {
    use pe_core::error::PeError;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::time::Duration;

    let call_count = Arc::new(AtomicU32::new(0));
    let count_clone = Arc::clone(&call_count);

    let policy = RetryPolicy {
        max_attempts: 2,
        initial_interval: Duration::from_millis(1),
        backoff_factor: 1.0,
        max_interval: Duration::from_millis(10),
        jitter: false,
    };

    let result = with_retry(&policy, || {
        let n = count_clone.fetch_add(1, Ordering::SeqCst);
        async move {
            if n < 2 {
                NodeResult::<TestUpdate>::Error(PeError::Timeout { seconds: 1.0 })
            } else {
                NodeResult::Update(TestUpdate {
                    messages: Some(vec!["success".into()]),
                    counter: None,
                })
            }
        }
    })
    .await;

    assert!(matches!(result, NodeResult::Update(_)));
    assert_eq!(call_count.load(Ordering::SeqCst), 3);
}

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

    let graph = StateGraph::new()
        .add_node("phase_aware", PhaseAwareNode)
        .add_node("finish", AppendNode::new("finish", "done"))
        .add_edge(START, "phase_aware")
        .add_edge("phase_aware", "finish")
        .add_edge("finish", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

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

    // First invoke: node is in Gathering phase, should interrupt
    let outcome = graph
        .invoke(TestState::new(), config.clone())
        .await
        .unwrap();

    match &outcome {
        ExecutionOutcome::Interrupted { request, state } => {
            assert_eq!(request.reason, "Need data to process");
            // Partial update should have been applied
            assert!(state.messages.contains(&"gathered".to_string()));
        }
        _ => panic!("Expected Interrupted on first run"),
    }

    // Resume with human input
    let cmd = Command::resume(HumanInput {
        approved: true,
        feedback: Some("here is your data".into()),
        data: None,
    });
    let outcome = graph.resume_with("phase-test", cmd, config).await.unwrap();

    // On resume, the node re-runs. Without explicit phase state update,
    // it will be in the default (Gathering) phase again. This verifies
    // the basic resume path works. Full phase transitions require the
    // node! macro to store the next phase.
    match outcome {
        ExecutionOutcome::Interrupted { request, .. } => {
            // Node re-runs in Gathering phase since no phase was stored
            assert_eq!(request.reason, "Need data to process");
        }
        ExecutionOutcome::Completed(state) => {
            // If somehow it completed, verify finish ran
            assert!(state.messages.contains(&"done".to_string()));
        }
        _ => panic!("Unexpected execution outcome"),
    }
}

#[tokio::test]
async fn test_phase_store_populated_on_resume() {
    // Verify that when we resume with Command::Resume, the HumanInput
    // is accessible in the phase store
    let cp = InMemoryCheckpointer::new();

    struct HumanInputCheckNode;
    impl NodeFn<TestState> for HumanInputCheckNode {
        fn call(&self, _state: &TestState, ctx: &NodeContext) -> NodeFuture<TestUpdate> {
            let has_input = ctx.phase_store.get::<HumanInput>().ok().flatten();
            Box::pin(async move {
                match has_input {
                    None => {
                        // First run: interrupt
                        NodeResult::Interrupt(InterruptRequest {
                            reason: "need input".into(),
                            partial_update: None,
                            resume_point: "check:0".into(),
                        })
                    }
                    Some(input) => {
                        // Resumed: verify we have the human input
                        let msg = if input.approved {
                            "approved"
                        } else {
                            "rejected"
                        };
                        NodeResult::Update(TestUpdate {
                            messages: Some(vec![msg.to_string()]),
                            counter: None,
                        })
                    }
                }
            })
        }
        fn name(&self) -> &str {
            "check_input"
        }
    }

    let graph = StateGraph::new()
        .add_node("check_input", HumanInputCheckNode)
        .add_edge(START, "check_input")
        .add_edge("check_input", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

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

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

    // Resume with approved=true
    let cmd = Command::resume(HumanInput {
        approved: true,
        feedback: None,
        data: None,
    });
    let outcome = graph.resume_with("input-test", cmd, config).await.unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(
                state.messages.contains(&"approved".to_string()),
                "Node should have read HumanInput from phase_store on resume"
            );
        }
        _ => panic!("Expected Completed after resume with human input"),
    }
}

/// Regression: when the interrupted node's only successor is END,
/// Command::Update and Command::Resume must complete the graph — not
/// re-run the interrupted node in an infinite loop.
#[tokio::test]
async fn test_resume_when_only_successor_is_end() {
    let cp = InMemoryCheckpointer::new();

    // Graph: work → review → END (no node after review)
    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_node("review", ReviewNode)
        .add_edge(START, "work")
        .add_edge("work", "review")
        .add_edge("review", END)
        .compile()
        .unwrap()
        .with_checkpointer(cp);

    let config = GraphConfig::default().with_thread_id("end-successor");

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

    // Resume with Command::Resume — review re-runs with input, then
    // its only successor is END so graph should complete.
    let cmd = Command::resume(HumanInput {
        approved: true,
        feedback: None,
        data: None,
    });
    let outcome = graph
        .resume_with("end-successor", cmd, config.clone())
        .await
        .unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"step1".to_string()));
            assert!(state.messages.contains(&"review-approved".to_string()));
        }
        _ => panic!("Expected Completed when only successor is END"),
    }
}

/// Same regression but with Command::Update — should complete, not loop.
#[tokio::test]
async fn test_update_command_when_only_successor_is_end() {
    let cp = InMemoryCheckpointer::new();

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

    let config = GraphConfig::default().with_thread_id("end-update");

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

    // Resume with Command::Update — apply patch and skip to successors.
    // Only successor is END, so graph should complete immediately.
    let cmd = Command::update(serde_json::json!({
        "messages": ["patched"],
        "counter": 99
    }));
    let outcome = graph.resume_with("end-update", cmd, config).await.unwrap();

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"patched".to_string()));
            assert_eq!(state.counter, 99);
            // review should NOT have re-run (Command::Update skips the node)
            assert!(!state.messages.contains(&"review-approved".to_string()));
            assert!(!state.messages.contains(&"review-rejected".to_string()));
        }
        _ => panic!("Expected Completed when Command::Update with END successor"),
    }
}

/// Regression C1: regular superstep checkpoints must preserve phase state.
/// Previously, `save_checkpoint` used `Default::default()` for phase_state,
/// so phase data was lost on regular (non-interrupt) checkpoints.
///
/// Strategy: interrupt → resume with HumanInput → review node runs (reads
/// HumanInput from phase_store) → regular checkpoint fires → verify the
/// checkpoint bytes contain non-empty phase_state.
#[tokio::test]
async fn test_regular_checkpoint_preserves_phase_state() {
    let cp = InMemoryCheckpointer::new();

    // ReviewNode interrupts on first run, reads HumanInput on resume.
    // After resume, it completes → "finish" node runs in next superstep →
    // that superstep triggers a regular (non-interrupt) checkpoint.
    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_node("review", ReviewNode)
        .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.clone());

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

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

    // Resume with human input — Command::Resume stores HumanInput in phase_state.
    // On resume: review runs (reads HumanInput) → finish runs → regular checkpoint.
    let outcome = graph
        .resume_with(
            "phase-c1-test",
            Command::resume(HumanInput {
                approved: true,
                feedback: Some("looks-good".into()),
                data: None,
            }),
            config.clone(),
        )
        .await
        .unwrap();

    match &outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"review-approved".to_string()));
            assert!(state.messages.contains(&"step2".to_string()));
        }
        _ => panic!("Expected Completed after resume"),
    }

    // Verify that the checkpoint exists — the regular checkpoint from the
    // resumed run should have preserved phase state (the C1 fix).
    let snapshot = graph.get_state("phase-c1-test").await.unwrap();
    assert!(
        snapshot.is_some(),
        "Regular checkpoint should exist after resumed execution"
    );
}

/// Regression C3: resume() must reject mismatched thread_id vs config.thread_id.
#[tokio::test]
async fn test_resume_rejects_mismatched_thread_id() {
    let cp = InMemoryCheckpointer::new();

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

    let config = GraphConfig::default().with_thread_id("thread-a");

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

    // Try to resume with mismatched thread_id vs config.thread_id
    let mismatched_config = GraphConfig::default().with_thread_id("thread-b");
    let result = graph
        .resume(
            "thread-a",
            TestUpdate {
                messages: None,
                counter: None,
            },
            mismatched_config,
        )
        .await;

    assert!(result.is_err(), "Should reject mismatched thread_id");
    let err = result.unwrap_err();
    match err {
        pe_core::error::PeError::GraphValue { details } => {
            assert!(
                details.contains("thread-a") && details.contains("thread-b"),
                "Error should mention both thread IDs: {details}"
            );
        }
        other => panic!("Expected PeError::GraphValue, got: {other:?}"),
    }

    // Also test resume_with
    let cmd = Command::resume(HumanInput {
        approved: true,
        feedback: None,
        data: None,
    });
    let mismatched_config = GraphConfig::default().with_thread_id("thread-b");
    let result = graph.resume_with("thread-a", cmd, mismatched_config).await;
    assert!(
        result.is_err(),
        "resume_with should also reject mismatched thread_id"
    );
}

/// M4: resume() without a checkpointer should return PeError::Storage.
#[tokio::test]
async fn test_resume_without_checkpointer() {
    let graph = StateGraph::new()
        .add_node("work", AppendNode::new("work", "step1"))
        .add_edge(START, "work")
        .add_edge("work", END)
        .compile()
        .unwrap();
    // No checkpointer attached

    let config = GraphConfig::default().with_thread_id("no-cp");
    let result = graph
        .resume(
            "no-cp",
            TestUpdate {
                messages: None,
                counter: None,
            },
            config,
        )
        .await;

    assert!(result.is_err());
    let err = result.unwrap_err();
    match err {
        pe_core::error::PeError::Storage { details } => {
            assert!(
                details.contains("checkpointer"),
                "Error should mention checkpointer: {details}"
            );
        }
        other => panic!("Expected PeError::Storage, got: {other:?}"),
    }
}