async-nodes 0.1.0

use std::{collections::HashSet, sync::Arc};

use async_nodes::{
    BspContext, BspEngine, BspError, BspGraph, BspInterceptor, BspOutputMapExt, BspType, BspValue,
    EngineResult, NodeId, NodeOutcome, OutputPin, node, node_blocking,
};
use indexmap::IndexMap;

// =========================================================================
// 1. TELEMETRY OBSERVER FOR TESTING
// =========================================================================

struct TestInterceptor {
    events: Arc<std::sync::Mutex<Vec<String>>>,
}

impl TestInterceptor {
    fn push_event(&self, event: String) {
        if let Ok(mut guard) = self.events.lock() {
            guard.push(event);
        }
    }
}

impl BspInterceptor for TestInterceptor {
    fn on_node_start(&self, _node_id: NodeId, name: &str) {
        self.push_event(format!("START: {name}"));
    }

    fn on_node_complete(&self, _node_id: NodeId, name: &str, _duration: std::time::Duration) {
        self.push_event(format!("COMPLETE: {name}"));
    }

    fn on_node_error(&self, _node_id: NodeId, name: &str, _err: &BspError) {
        self.push_event(format!("ERROR: {name}"));
    }
}

// =========================================================================
// 2. CONTEXT STRUCTS FOR DEPENDENCY INJECTION
// =========================================================================

#[derive(Clone, Debug)]
struct DatabaseConnection {
    url: String,
}

impl async_nodes::typemap_rev::TypeMapKey for DatabaseConnection {
    type Value = Self;
}

#[derive(Clone, Debug)]
struct AppConfig {
    timeout_ms: u64,
}

impl async_nodes::typemap_rev::TypeMapKey for AppConfig {
    type Value = Self;
}

#[derive(Clone, Debug)]
struct HeavyResource {
    data: Vec<u8>,
}

impl async_nodes::typemap_rev::TypeMapKey for HeavyResource {
    type Value = std::sync::Arc<Self>;
}

// =========================================================================
// 3. TEST NODE DEFINITIONS USING #[node] AND #[node_blocking]
// =========================================================================

#[node("Generator")]
async fn generate_value() -> f64 {
    10.0
}

#[node("Multiply")]
async fn multiply_by_three(input: f64) -> f64 {
    input * 3.0
}

#[node("Source")]
async fn source_value() -> f64 {
    5.0
}

#[node("AddTen")]
async fn add_ten(val: f64) -> f64 {
    tokio::task::yield_now().await;
    val + 10.0
}

#[node("Double")]
async fn double_value(val: f64) -> f64 {
    tokio::task::yield_now().await;
    val * 2.0
}

#[node("Join")]
async fn join_values(left: f64, right: f64) -> f64 {
    left + right
}

#[node_blocking("ComputeHeavy")]
fn run_heavy_blocking_task(input: f64) -> f64 {
    // Simulate a synchronous thread block
    std::thread::sleep(std::time::Duration::from_millis(150));
    input * 5.0
}

#[node("Suspender")]
async fn suspension_checkpoint(val: f64, approved: Option<bool>) -> Result<NodeOutcome, String> {
    if approved == Some(true) {
        let mut out = IndexMap::new();
        out.insert(0, BspValue::new(val));
        out.insert(1, BspValue::new(true));
        Ok(NodeOutcome::Success(out))
    } else if val > 1.0 {
        Ok(NodeOutcome::Suspend {
            payload: serde_json::json!({ "prompt": "Authorize payout?", "amount": val }),
            waiting_inputs: vec![1],
        })
    } else {
        let mut out = IndexMap::new();
        out.insert(0, BspValue::new(val));
        out.insert(1, BspValue::new(false));
        Ok(NodeOutcome::Success(out))
    }
}

#[node("PostAuth")]
async fn handle_post_auth(original_amount: f64, approved: bool) -> f64 {
    if approved { original_amount } else { 0.0 }
}

// Custom structure utilizing the structural type-derive
#[derive(Clone, Debug, PartialEq, BspType)]
struct Transaction {
    id: String,
    amount: f64,
}

#[node("CreateTransaction")]
async fn create_transaction() -> Transaction {
    Transaction {
        id: "tx-4509".to_string(),
        amount: 150.0,
    }
}

#[node("CalculateTaxes")]
async fn calculate_taxes(tx: Transaction) -> f64 {
    tx.amount * 0.15
}

#[node("FormatFloat")]
async fn format_float(val: f64) -> String {
    format!("{val:.1}")
}

#[node("JoinStringCollection")]
async fn join_string_collection(input: Vec<String>) -> String {
    input.join(", ")
}

// Nodes leveraging context injection via owned-cloned semantics
#[node("QueryUser")]
async fn query_user(#[context] db: DatabaseConnection, user_id: f64) -> String {
    format!("User 42 from {} with id {}", db.url, user_id)
}

#[node("FetchTimeout")]
async fn fetch_timeout(#[context] config: AppConfig) -> f64 {
    #[allow(clippy::cast_precision_loss)]
    let timeout = config.timeout_ms as f64;

    timeout
}

// Blocking Node utilizing context-injected Arc (mapped explicitly to TypeMapKey HeavyResource)
#[node_blocking("HeavyDiskWrite")]
fn heavy_disk_write(
    #[context(HeavyResource)] resource: std::sync::Arc<HeavyResource>,
    file_prefix: String,
) -> String {
    std::thread::sleep(std::time::Duration::from_millis(50));
    format!(
        "Wrote prefix '{}' with resource size: {}",
        file_prefix,
        resource.data.len()
    )
}

// =========================================================================
// 4. CORE PIPELINE TESTS
// =========================================================================

#[tokio::test]
async fn test_sequential_pipeline() {
    let mut graph = BspGraph::new();

    let gen_node = graph.add_node(GenerateValueNode::new());
    let mul = graph.add_node(MultiplyByThreeNode::new());

    // Wire Generator Output (0) -> Multiply Input (0)
    graph.add_link(gen_node.output(), mul.input());

    let engine = BspEngine::new(graph);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let final_value: f64 = outputs.get_typed(mul.output()).unwrap();
        assert!((final_value - 30.0).abs() < f64::EPSILON);
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_blocking_node_execution() {
    let mut graph = BspGraph::new();

    let src = graph.add_node(SourceValueNode::new()); // Outputs 5.0
    let heavy = graph.add_node(RunHeavyBlockingTaskNode::new()); // 5.0 * 5.0 = 25.0

    graph.add_link(src.output(), heavy.input());

    let engine = BspEngine::new(graph);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let final_value: f64 = outputs.get_typed(heavy.output()).unwrap();
        assert!((final_value - 25.0).abs() < f64::EPSILON);
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_parallel_branch_join_determinism() {
    let mut graph = BspGraph::new();

    // Node 1: Source (Outputs 5.0)
    let src = graph.add_node(SourceValueNode::new());

    // Node 2 & 3: Level 1 Parallel Branches (AddTen & Double)
    let add = graph.add_node(AddTenNode::new());
    let dbl = graph.add_node(DoubleValueNode::new());

    // Node 4: Level 2 Join (Adds parallel outcomes together)
    let join = graph.add_node(JoinValuesNode::new());

    // Link output of Source (1) to inputs of AddTen (2) and Double (3)
    graph.add_link(src.output(), add.val());
    graph.add_link(src.output(), dbl.val());

    // Link outputs of AddTen (2) and Double (3) to inputs 0 and 1 of Join (4)
    graph.add_link(add.output(), join.left());
    graph.add_link(dbl.output(), join.right());

    let events = Arc::new(std::sync::Mutex::new(Vec::new()));
    let interceptor = Arc::new(TestInterceptor {
        events: events.clone(),
    });

    let engine = BspEngine::new(graph).with_interceptor(interceptor);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let sum: f64 = outputs.get_typed(join.output()).unwrap();

        // Mathematical verification: (5.0 + 10.0) + (5.0 * 2.0) = 25.0
        assert!((sum - 25.0).abs() < f64::EPSILON);

        // Verification of concurrency logging via interceptor
        // Assert Source starts first
        let events = events.lock().unwrap().clone();
        assert_eq!(events[0], "START: Source");
        assert_eq!(events[1], "COMPLETE: Source");

        // Assert Level 1 nodes ran in the same superstep
        let level_1_started: HashSet<&str> = events[2..4].iter().map(String::as_str).collect();
        assert!(level_1_started.contains("START: AddTen"));
        assert!(level_1_started.contains("START: Double"));

        // Assert Join Node was executed strictly last
        assert_eq!(*events.last().unwrap(), "COMPLETE: Join");
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_graph_suspension_and_resumption() {
    let mut graph = BspGraph::new();
    // Register nodes
    let src = graph.add_node(SourceValueNode::new()); // Outputs 5.0
    let susp = graph.add_node(SuspensionCheckpointNode::new()); // Suspends if input > 1.0
    let auth = graph.add_node(HandlePostAuthNode::new()); // Resolves outcome

    // Wire: Source (1) -> Suspender (2)
    graph.add_link(src.output(), susp.val());

    // Explicit dynamic wire mapping
    // Wire: Suspender (2) -> PostAuth (3)
    graph.add_link(
        OutputPin {
            node_id: susp.id,
            port: 0,
        },
        auth.original_amount(),
    );
    graph.add_link(
        OutputPin {
            node_id: susp.id,
            port: 1,
        },
        auth.approved(),
    );

    let engine = BspEngine::new(graph);
    let result = engine.run(()).await.unwrap();

    match result {
        EngineResult::Suspended(suspended) => {
            // Verify suspension telemetry payload matches expectation
            assert_eq!(suspended.suspensions.len(), 1);
            let suspension = &suspended.suspensions[0];
            assert_eq!(suspension.node_id, susp.id);
            assert_eq!(suspension.waiting_inputs, vec![1]);

            // Simulate host intervention: Inject authorization boolean (true) into waiting input pin (1) of node (2)
            let mut injections = IndexMap::new();
            injections.insert(susp.approved(), BspValue::new(true));

            // Resume execution
            let resumed_result = engine.resume(*suspended, injections).await.unwrap();

            if let EngineResult::Complete(outputs) = resumed_result {
                let amount: f64 = outputs.get_typed(auth.output()).unwrap();

                // PostAuth output should match approved original amount
                assert!((amount - 5.0).abs() < f64::EPSILON);
            } else {
                panic!("Resumed execution failed to complete.");
            }
        }
        EngineResult::Complete(_) => panic!("Graph execution should have suspended."),
        EngineResult::StepLimitReached(_) => panic!("Graph execution should have suspended."),
    }
}

#[tokio::test]
async fn test_cycle_validation() {
    let mut graph = BspGraph::new();

    // Create an invalid cyclic loop: Node 1 -> Node 2 -> Node 1
    let node1 = graph.add_node(AddTenNode::new());
    let node2 = graph.add_node(DoubleValueNode::new());

    graph.add_link(node1.output(), node2.val());
    graph.add_link(node2.output(), node1.val());

    let engine = BspEngine::new(graph);
    let result = engine.run(()).await;

    // Verify engine caught circular deadlock before executing
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        BspError::CycleDetected { .. }
    ));
}

#[tokio::test]
async fn test_custom_type_and_ergonomics() {
    let mut graph = BspGraph::new();

    // 1. Auto-incrementing node additions
    let src = graph.add_node(CreateTransactionNode::new());
    let tax = graph.add_node(CalculateTaxesNode::new());

    // 2. Type-safe, handle-based linking
    graph.add_link(src.output(), tax.tx());

    let engine = BspEngine::new(graph);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        // 3. Clean type-safe downcasting helper
        let calculated_tax: f64 = outputs.get_typed(tax.output()).unwrap();
        assert!((calculated_tax - 22.5).abs() < f64::EPSILON);
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_multiple_inputs_into_single_port() {
    let mut graph = BspGraph::new();

    // Source 1 and Source 2
    let src1 = graph.add_node(AddTenNode::new()); // Output is f64
    let src2 = graph.add_node(DoubleValueNode::new()); // Output is f64

    // Formatter nodes to turn floats into Strings
    let fmt1 = graph.add_node(FormatFloatNode::new());
    let fmt2 = graph.add_node(FormatFloatNode::new());

    // Target aggregation node taking Vec<String>
    let joiner = graph.add_node(JoinStringCollectionNode::new());

    // Set static inputs to our parallel pipelines
    graph.set_static_input(src1.val(), 10.0); // 10.0 + 10.0 = 20.0
    graph.set_static_input(src2.val(), 10.0); // 10.0 * 2.0 = 20.0

    // Connect source nodes to formatters
    graph.add_link(src1.output(), fmt1.val());
    graph.add_link(src2.output(), fmt2.val());

    // Connect BOTH formatters into the exact same input port of the joiner
    graph.add_link(fmt1.output(), joiner.input());
    graph.add_link(fmt2.output(), joiner.input());

    let engine = BspEngine::new(graph);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let final_string: String = outputs.get_typed(joiner.output()).unwrap();

        // Assert both strings are compiled, aggregated, and joined in parallel supersteps
        let parts: Vec<&str> = final_string.split(", ").collect();
        assert_eq!(parts.len(), 2);
        assert!(parts.contains(&"20.0"));
    } else {
        panic!("The execution pipeline should have finished successfully.");
    }
}

// =========================================================================
// 5. CONTEXT DEPENDENCY INJECTION TESTS
// =========================================================================

#[tokio::test]
async fn test_dependency_injection_async_owned() {
    let mut graph = BspGraph::new();

    let query_node = graph.add_node(QueryUserNode::new());
    graph.set_static_input(query_node.user_id(), 42.0);

    let ctx = BspContext::new();
    ctx.insert::<DatabaseConnection>(DatabaseConnection {
        url: "postgres://localhost:5432".to_string(),
    });

    let engine = BspEngine::new(graph).with_context(ctx);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let outcome: String = outputs.get_typed(query_node.output()).unwrap();
        assert_eq!(outcome, "User 42 from postgres://localhost:5432 with id 42");
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_dependency_injection_async_clone() {
    let mut graph = BspGraph::new();

    let config_node = graph.add_node(FetchTimeoutNode::new());

    let ctx = BspContext::new();
    ctx.insert::<AppConfig>(AppConfig { timeout_ms: 1500 });

    let engine = BspEngine::new(graph).with_context(ctx);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let val: f64 = outputs.get_typed(config_node.output()).unwrap();
        assert!((val - 1500.0).abs() < f64::EPSILON);
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_dependency_injection_blocking_task() {
    let mut graph = BspGraph::new();

    let write_node = graph.add_node(HeavyDiskWriteNode::new());
    graph.set_static_input(write_node.file_prefix(), "backup_v1".to_string());

    let ctx = BspContext::new();
    ctx.insert::<HeavyResource>(std::sync::Arc::new(HeavyResource {
        data: vec![0, 1, 2, 3, 4],
    }));

    let engine = BspEngine::new(graph).with_context(ctx);
    let result = engine.run(()).await.unwrap();

    if let EngineResult::Complete(outputs) = result {
        let outcome: String = outputs.get_typed(write_node.output()).unwrap();
        assert_eq!(outcome, "Wrote prefix 'backup_v1' with resource size: 5");
    } else {
        panic!("Execution should have finished to completion.");
    }
}

#[tokio::test]
async fn test_missing_context_dependency_validation_error() {
    let mut graph = BspGraph::new();
    let _query_node = graph.add_node(QueryUserNode::new());

    let ctx = BspContext::new();
    let engine = BspEngine::new(graph).with_context(ctx);
    let result = engine.run(()).await;

    assert!(result.is_err());
    let err_msg = result.unwrap_err().to_string();
    assert!(err_msg.contains("Dependency injection failed"));
    assert!(err_msg.contains("DatabaseConnection"));
}