pe-graph 0.1.0

//! Integration tests for the matrix layer wiring into the Pregel engine.
//!
//! Tests that convergence signals are recorded, C value updates, matrix-guided
//! routing works, and non-matrix mode degrades cleanly.

use pe_core::node::{ConvergenceSignal, NodeContext, NodeFn, NodeFuture, NodeResult};
use pe_core::state::{State, StateUpdate};
use pe_graph::matrix_hook::{
    ConvergenceRecorder, DefaultMatrixHook, MatrixHookHandle, RoutingResolver,
};
use pe_graph::{END, ExecutionOutcome, GraphConfig, START, StateGraph};
use serde::{Deserialize, Serialize};

// ── Test State ──────────────────────────────────────────────────────

#[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 Tracker/Router (implement pe-graph traits directly) ────────

#[derive(Debug, Clone)]
struct TestTracker {
    c_value: f64,
    completion: f64,
    signal_count: u32,
    threshold: f64,
    observations: Vec<(f64, f64, f64)>,
}

impl TestTracker {
    fn new(threshold: f64) -> Self {
        Self {
            c_value: 0.0,
            completion: 0.0,
            signal_count: 0,
            threshold,
            observations: Vec::new(),
        }
    }
}

impl ConvergenceRecorder for TestTracker {
    fn record(&mut self, contribution: f64, surprise: f64, quality: f64) {
        self.signal_count += 1;
        self.observations.push((contribution, surprise, quality));
        let adjusted = quality * (1.0 - surprise);
        self.c_value = self.c_value * 0.5 + adjusted * 0.5;
        self.completion = (self.completion + contribution).min(1.0);
    }

    fn c_value(&self) -> f64 {
        self.c_value
    }

    fn completion(&self) -> f64 {
        self.completion
    }

    fn is_converged(&self) -> bool {
        self.c_value >= self.threshold
    }
}

struct TestRouter {
    overrides: std::collections::HashMap<String, String>,
    transitions: Vec<(String, String, f64)>,
}

impl TestRouter {
    fn new() -> Self {
        Self {
            overrides: std::collections::HashMap::new(),
            transitions: Vec::new(),
        }
    }

    fn override_route(&mut self, from: &str, to: &str) {
        self.overrides.insert(from.to_string(), to.to_string());
    }
}

impl RoutingResolver for TestRouter {
    fn resolve(&self, from: &str, candidates: &[String]) -> Option<Vec<String>> {
        if let Some(target) = self.overrides.get(from) {
            if candidates.contains(target) {
                return Some(vec![target.clone()]);
            }
        }
        None // Fall back to user's router
    }

    fn learn(&mut self, from: &str, to: &str, quality: f64) {
        self.transitions
            .push((from.to_string(), to.to_string(), quality));
    }
}

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

/// Node that returns a convergence signal.
struct ConvergeNode {
    name: &'static str,
    contribution: f64,
    surprise: f64,
    quality: f64,
    message: &'static str,
}

impl NodeFn<TestState> for ConvergeNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let signal = ConvergenceSignal {
            actual_contribution: self.contribution,
            surprise: self.surprise,
            quality: self.quality,
            partial_update: TestUpdate {
                messages: Some(vec![self.message.to_string()]),
                counter: None,
            },
        };
        Box::pin(async move { NodeResult::Converge(signal) })
    }

    fn name(&self) -> &str {
        self.name
    }
}

/// Simple append node.
struct AppendNode {
    name: &'static str,
    message: &'static str,
}

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.name
    }
}

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

#[tokio::test]
async fn converge_signal_updates_c_value_via_hook() {
    let tracker = TestTracker::new(0.95);
    let router = TestRouter::new();
    let hook = MatrixHookHandle::new(DefaultMatrixHook::new(tracker, router));

    let graph = StateGraph::new()
        .add_node(
            "converge",
            ConvergeNode {
                name: "converge",
                contribution: 0.5,
                surprise: 0.0,
                quality: 0.9,
                message: "converged",
            },
        )
        .add_edge(START, "converge")
        .add_edge("converge", END)
        .compile()
        .unwrap()
        .with_matrix_hook(hook.clone());

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

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

    // The hook should have recorded the convergence signal
    assert!(
        hook.c_value() > 0.0,
        "c_value should be > 0 after convergence signal"
    );
    assert!(
        (hook.completion() - 0.5).abs() < 0.001,
        "completion should be ~0.5"
    );
}

#[tokio::test]
async fn multiple_converge_signals_increase_completion() {
    let tracker = TestTracker::new(0.95);
    let router = TestRouter::new();
    let hook = MatrixHookHandle::new(DefaultMatrixHook::new(tracker, router));

    // Build a 3-node chain where each returns a convergence signal
    let graph = StateGraph::new()
        .add_node(
            "step1",
            ConvergeNode {
                name: "step1",
                contribution: 0.3,
                surprise: 0.0,
                quality: 0.8,
                message: "step1",
            },
        )
        .add_node(
            "step2",
            ConvergeNode {
                name: "step2",
                contribution: 0.3,
                surprise: 0.0,
                quality: 0.9,
                message: "step2",
            },
        )
        .add_node(
            "step3",
            ConvergeNode {
                name: "step3",
                contribution: 0.4,
                surprise: 0.0,
                quality: 1.0,
                message: "step3",
            },
        )
        .add_edge(START, "step1")
        .add_edge("step1", "step2")
        .add_edge("step2", "step3")
        .add_edge("step3", END)
        .compile()
        .unwrap()
        .with_matrix_hook(hook.clone());

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

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages.len(), 3);
        }
        _ => panic!("expected Completed"),
    }

    // Completion should be 0.3 + 0.3 + 0.4 = 1.0
    assert!(
        (hook.completion() - 1.0).abs() < 0.001,
        "completion should approach 1.0, got {}",
        hook.completion()
    );
    assert!(hook.c_value() > 0.0, "c_value should be > 0");
}

#[tokio::test]
async fn converge_without_matrix_degrades_to_update() {
    // No matrix hook attached — Converge should still apply partial_update
    let graph = StateGraph::new()
        .add_node(
            "converge",
            ConvergeNode {
                name: "converge",
                contribution: 0.5,
                surprise: 0.0,
                quality: 0.9,
                message: "still applied",
            },
        )
        .add_edge(START, "converge")
        .add_edge("converge", END)
        .compile()
        .unwrap();
    // No with_matrix_hook — this is the degradation path

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

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

#[tokio::test]
async fn matrix_router_overrides_conditional_edge() {
    let tracker = TestTracker::new(0.95);
    let mut router = TestRouter::new();
    // Override: when conditional edge from "decide" fires, always go to "path_b"
    router.override_route("decide", "path_b");
    let hook = MatrixHookHandle::new(DefaultMatrixHook::new(tracker, router));

    let graph = StateGraph::new()
        .add_node(
            "decide",
            AppendNode {
                name: "decide",
                message: "decided",
            },
        )
        .add_node(
            "path_a",
            AppendNode {
                name: "path_a",
                message: "took_a",
            },
        )
        .add_node(
            "path_b",
            AppendNode {
                name: "path_b",
                message: "took_b",
            },
        )
        .add_edge(START, "decide")
        // Conditional edge: user's router returns BOTH paths
        .add_conditional_edge("decide", |_: &TestState| {
            vec!["path_a".to_string(), "path_b".to_string()]
        })
        .add_edge("path_a", END)
        .add_edge("path_b", END)
        .compile()
        .unwrap()
        .with_matrix_hook(hook);

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

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"decided".to_string()));
            // Matrix router should have overridden to only path_b
            assert!(state.messages.contains(&"took_b".to_string()));
            assert!(!state.messages.contains(&"took_a".to_string()));
        }
        _ => panic!("expected Completed"),
    }
}

#[tokio::test]
async fn conditional_edge_works_without_matrix() {
    // No matrix hook — conditional edge uses user's router directly
    let graph = StateGraph::new()
        .add_node(
            "decide",
            AppendNode {
                name: "decide",
                message: "decided",
            },
        )
        .add_node(
            "path_a",
            AppendNode {
                name: "path_a",
                message: "took_a",
            },
        )
        .add_node(
            "path_b",
            AppendNode {
                name: "path_b",
                message: "took_b",
            },
        )
        .add_edge(START, "decide")
        // User's router sends to path_a only
        .add_conditional_edge("decide", |_: &TestState| vec!["path_a".to_string()])
        .add_edge("path_a", END)
        .add_edge("path_b", END)
        .compile()
        .unwrap();

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

    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert!(state.messages.contains(&"took_a".to_string()));
            assert!(!state.messages.contains(&"took_b".to_string()));
        }
        _ => panic!("expected Completed"),
    }
}

#[tokio::test]
async fn matrix_router_fallback_when_no_override() {
    let tracker = TestTracker::new(0.95);
    let router = TestRouter::new(); // No overrides — will fall back
    let hook = MatrixHookHandle::new(DefaultMatrixHook::new(tracker, router));

    let graph = StateGraph::new()
        .add_node(
            "decide",
            AppendNode {
                name: "decide",
                message: "decided",
            },
        )
        .add_node(
            "path_a",
            AppendNode {
                name: "path_a",
                message: "took_a",
            },
        )
        .add_edge(START, "decide")
        .add_conditional_edge("decide", |_: &TestState| vec!["path_a".to_string()])
        .add_edge("path_a", END)
        .compile()
        .unwrap()
        .with_matrix_hook(hook);

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

    match outcome {
        ExecutionOutcome::Completed(state) => {
            // Should have fallen back to user's router
            assert!(state.messages.contains(&"took_a".to_string()));
        }
        _ => panic!("expected Completed"),
    }
}

#[tokio::test]
async fn transitions_recorded_for_learning() {
    // Use shared state so we can observe transitions after execution.
    let transition_log =
        std::sync::Arc::new(std::sync::Mutex::new(Vec::<(String, String, f64)>::new()));

    struct LoggingRouter {
        log: std::sync::Arc<std::sync::Mutex<Vec<(String, String, f64)>>>,
    }
    impl RoutingResolver for LoggingRouter {
        fn resolve(&self, _from: &str, _candidates: &[String]) -> Option<Vec<String>> {
            None // fall back to user's router
        }
        fn learn(&mut self, from: &str, to: &str, quality: f64) {
            self.log
                .lock()
                .unwrap()
                .push((from.into(), to.into(), quality));
        }
    }

    let tracker = TestTracker::new(0.95);
    let router = LoggingRouter {
        log: transition_log.clone(),
    };
    let hook = MatrixHookHandle::new(DefaultMatrixHook::new(tracker, router));

    // Use ConvergeNode so transitions are actually recorded
    // (only Converge signals trigger record_transition)
    let graph = StateGraph::new()
        .add_node(
            "a",
            ConvergeNode {
                name: "a",
                contribution: 0.3,
                surprise: 0.0,
                quality: 0.8,
                message: "a",
            },
        )
        .add_node(
            "b",
            ConvergeNode {
                name: "b",
                contribution: 0.3,
                surprise: 0.0,
                quality: 0.9,
                message: "b",
            },
        )
        .add_edge(START, "a")
        .add_edge("a", "b")
        .add_edge("b", END)
        .compile()
        .unwrap()
        .with_matrix_hook(hook.clone());

    let outcome = graph
        .invoke(TestState::new(), GraphConfig::default())
        .await
        .unwrap();
    assert!(matches!(outcome, ExecutionOutcome::Completed(_)));

    // Verify transitions were actually recorded with correct quality values
    let log = transition_log.lock().unwrap();
    assert!(
        !log.is_empty(),
        "transitions should be recorded for Converge nodes"
    );
    // a->b transition with quality 0.8 (from node "a"'s Converge signal)
    assert!(
        log.iter().any(|(from, to, _)| from == "a" && to == "b"),
        "expected a->b transition, got: {:?}",
        *log
    );
}