resuma 1.3.1

//! Flow engine — start workers, hold graph executions, pause/resume hooks.

use std::collections::HashMap;
use std::sync::Arc;

use once_cell::sync::Lazy;
use parking_lot::RwLock;
use serde_json::Value;
use tokio::task;
use tokio_util::sync::CancellationToken;

use crate::core::{Result, ResumaError};

use super::cancel;
use super::durable::{self, ExecutionRecord};
use super::events::{emit, EventBus, SharedEventBus};
use super::graph;
use super::planner::{self, PlannerHints};
use super::resources::{self, ResourceProfile};
use super::runner;
use super::state::StateStore;
use super::types::{GraphId, GraphSnapshot, GraphStatus, StartWorkerResponse};
use super::workers::{self, emit_worker_start, WorkerContext, WorkerFn};

static GRAPHS: Lazy<RwLock<HashMap<String, Arc<GraphExecution>>>> =
    Lazy::new(|| RwLock::new(HashMap::new()));

static RESUME_LOCKS: Lazy<RwLock<HashMap<String, Arc<tokio::sync::Mutex<()>>>>> =
    Lazy::new(|| RwLock::new(HashMap::new()));

fn resume_lock(id: &str) -> Arc<tokio::sync::Mutex<()>> {
    let mut map = RESUME_LOCKS.write();
    map.entry(id.to_string())
        .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
        .clone()
}

/// Drop in-memory execution state once a graph reaches a terminal status.
/// Snapshots remain available via durable storage.
fn release_live_graph(id: &GraphId) {
    GRAPHS.write().remove(&id.0);
}

/// Live graph execution state.
pub struct GraphExecution {
    pub snapshot: Arc<RwLock<GraphSnapshot>>,
    pub bus: SharedEventBus,
    pub state: Arc<StateStore>,
    pub worker: String,
    pub input: Value,
    pub profile: ResourceProfile,
    pub plan: super::types::ExecutionPlan,
    pub cancel: CancellationToken,
}

/// Start a registered worker; returns immediately while execution runs in background.
pub struct FlowEngine;

/// Live graph counts for ops / status.
#[derive(Debug, Clone, Copy, Default)]
pub struct GraphCounts {
    pub active: usize,
    pub running: usize,
    pub paused: usize,
}

impl FlowEngine {
    pub fn graph_counts() -> GraphCounts {
        let graphs = GRAPHS.read();
        let mut running = 0usize;
        let mut paused = 0usize;
        for g in graphs.values() {
            match g.snapshot.read().status {
                GraphStatus::Running => running += 1,
                GraphStatus::Paused => paused += 1,
                _ => {}
            }
        }
        GraphCounts {
            active: graphs.len(),
            running,
            paused,
        }
    }

    pub async fn start(name: &str, input: Value) -> Result<StartWorkerResponse> {
        let (meta, run) = workers::get_worker(name)
            .ok_or_else(|| ResumaError::UnknownWorker(name.to_string()))?;

        let hints = PlannerHints {
            use_ai: meta.intent.to_lowercase().contains("ai"),
            tools: Vec::new(),
        };
        let plan = planner::plan(&meta.intent, hints);
        let profile = resources::resolve(&meta.resources, &plan);

        let graph_id = GraphId::new();
        let access_token = super::security::issue_graph_token(&graph_id)?;
        super::metrics::inc_graph_started();
        spawn_execution(
            graph_id.clone(),
            name.to_string(),
            meta.intent.clone(),
            input,
            plan.clone(),
            profile.clone(),
            run,
            None,
            None,
            None,
        )
        .await?;

        Ok(StartWorkerResponse {
            graph_id,
            plan,
            access_token: Some(access_token),
        })
    }

    /// Resume a paused graph from durable checkpoint.
    pub async fn resume(id: &GraphId) -> Result<StartWorkerResponse> {
        let lock = resume_lock(&id.0);
        let _guard = lock.lock().await;

        let record = durable::load_execution_record(id)
            .ok_or_else(|| ResumaError::UnknownGraph(id.0.clone()))?;
        if record.cancelled {
            return Err(ResumaError::validation("graph was cancelled"));
        }
        if !record.paused {
            return Err(ResumaError::validation("graph is not paused"));
        }

        // Reuse the live EventBus so open SSE streams keep receiving events
        // after Resume (pause cancels the worker but the browser stays subscribed).
        let reuse_bus = GRAPHS.read().get(&id.0).map(|e| e.bus.clone());

        if let Some(exec) = GRAPHS.read().get(&id.0) {
            if exec.snapshot.read().status == GraphStatus::Running {
                return Err(ResumaError::validation("graph already running"));
            }
        }

        let (_, run) = workers::get_worker(&record.worker)
            .ok_or_else(|| ResumaError::UnknownWorker(record.worker.clone()))?;

        let snapshot =
            durable::load_graph(id).ok_or_else(|| ResumaError::UnknownGraph(id.0.clone()))?;
        let state = durable::load_checkpoint(id).unwrap_or_default();

        spawn_execution(
            id.clone(),
            record.worker.clone(),
            snapshot.intent.clone(),
            record.input.clone(),
            record.plan.clone(),
            record.profile.clone(),
            run,
            Some(snapshot),
            Some(state),
            reuse_bus,
        )
        .await?;

        let mut record = record;
        record.paused = false;
        let _ = durable::save_execution_record(&record);

        Ok(StartWorkerResponse {
            graph_id: id.clone(),
            plan: record.plan,
            access_token: durable::load_graph_token(id),
        })
    }

    pub fn snapshot(id: &GraphId) -> Option<GraphSnapshot> {
        GRAPHS
            .read()
            .get(&id.0)
            .map(|g| g.snapshot.read().clone())
            .or_else(|| durable::load_graph(id))
    }

    pub fn bus(id: &GraphId) -> Option<SharedEventBus> {
        GRAPHS.read().get(&id.0).map(|g| g.bus.clone())
    }

    pub fn replay(id: &GraphId) -> Option<Vec<super::types::WorkerEvent>> {
        GRAPHS
            .read()
            .get(&id.0)
            .map(|g| g.bus.history())
            .or_else(|| durable::load_events(id))
    }

    /// The `artifact_id` from the most recent `WorkerEvent::Result` in this
    /// graph's event history, if any.
    ///
    /// Workers that return large payloads via [`crate::exec::artifact_put`] put
    /// the resulting id in their result JSON under `"artifact_id"` by convention;
    /// callers otherwise had to `replay` the whole event history and search it by
    /// hand on every claim/poll endpoint. Returns `None` if the graph has no
    /// history, or its last result carries no `artifact_id` field.
    pub fn last_artifact(id: &GraphId) -> Option<String> {
        Self::last_result_field(id, "artifact_id")
    }

    /// The value of `field` in the most recent `WorkerEvent::Result` payload for
    /// this graph, as a string. General form of [`Self::last_artifact`] for
    /// workers that key their result differently (e.g. `"chunk_id"`).
    ///
    /// JSON strings, numbers, and booleans are coerced to `String`; objects,
    /// arrays, and null yield `None`.
    pub fn last_result_field(id: &GraphId, field: &str) -> Option<String> {
        let events = Self::replay(id)?;
        extract_last_result_field(&events, field)
    }

    /// Pause and **cancel** the in-flight worker (cooperative abort, resumable).
    pub fn pause(id: &GraphId) -> Result<()> {
        let exec = GRAPHS
            .read()
            .get(&id.0)
            .cloned()
            .or_else(|| restore_exec_from_durable(id))
            .ok_or_else(|| ResumaError::UnknownGraph(id.0.clone()))?;

        {
            let snap = exec.snapshot.read();
            match snap.status {
                GraphStatus::Running => {}
                GraphStatus::Paused => return Ok(()),
                GraphStatus::Done | GraphStatus::Failed => {
                    return Err(ResumaError::validation("cannot pause finished graph"));
                }
            }
        }

        // Signal cancellation first so run_on_node / map-reduce stop promptly.
        exec.cancel.cancel();

        {
            let mut snap = exec.snapshot.write();
            graph::mark_paused(&mut snap);
            let _ = durable::persist_graph(&snap);
        }
        let _ = durable::save_checkpoint(id, &exec.state);
        let _ = durable::persist_events(id, &exec.bus.history());

        let record = ExecutionRecord {
            graph_id: id.clone(),
            worker: exec.worker.clone(),
            input: exec.input.clone(),
            plan: exec.plan.clone(),
            profile: exec.profile.clone(),
            paused: true,
            cancelled: false,
        };
        let _ = durable::save_execution_record(&record);

        GRAPHS.write().insert(id.0.clone(), exec);
        Ok(())
    }

    /// Cancel a graph permanently (not resumable). Running workers are aborted.
    pub fn cancel(id: &GraphId) -> Result<()> {
        let exec = GRAPHS
            .read()
            .get(&id.0)
            .cloned()
            .or_else(|| restore_exec_from_durable(id))
            .ok_or_else(|| ResumaError::UnknownGraph(id.0.clone()))?;

        {
            let snap = exec.snapshot.read();
            match snap.status {
                GraphStatus::Done | GraphStatus::Failed => {
                    return Err(ResumaError::validation("cannot cancel finished graph"));
                }
                GraphStatus::Paused => {
                    let mut snap = exec.snapshot.write();
                    graph::mark_failed(&mut snap);
                    let _ = durable::persist_graph(&snap);
                    let record = ExecutionRecord {
                        graph_id: id.clone(),
                        worker: exec.worker.clone(),
                        input: exec.input.clone(),
                        plan: exec.plan.clone(),
                        profile: exec.profile.clone(),
                        paused: false,
                        cancelled: true,
                    };
                    let _ = durable::save_execution_record(&record);
                    super::metrics::inc_graph_failed();
                    super::webhooks::notify_failed(&snap, 0, "cancelled by operator".into());
                    return Ok(());
                }
                GraphStatus::Running => {}
            }
        }

        exec.cancel.cancel();

        let record = ExecutionRecord {
            graph_id: id.clone(),
            worker: exec.worker.clone(),
            input: exec.input.clone(),
            plan: exec.plan.clone(),
            profile: exec.profile.clone(),
            paused: false,
            cancelled: true,
        };
        let _ = durable::save_execution_record(&record);

        GRAPHS.write().insert(id.0.clone(), exec);
        Ok(())
    }
}

#[allow(clippy::too_many_arguments)]
async fn spawn_execution(
    graph_id: GraphId,
    worker_name: String,
    intent: String,
    input: Value,
    plan: super::types::ExecutionPlan,
    profile: ResourceProfile,
    run: WorkerFn,
    existing_snapshot: Option<GraphSnapshot>,
    existing_state: Option<StateStore>,
    reuse_bus: Option<SharedEventBus>,
) -> Result<()> {
    let snapshot = existing_snapshot
        .unwrap_or_else(|| graph::materialize(graph_id.clone(), &worker_name, &intent, &plan));
    let reused_bus = reuse_bus.is_some();
    let bus = reuse_bus.unwrap_or_else(|| Arc::new(EventBus::new()));

    // Only seed history on a fresh bus — reused buses already carry live + prior events.
    if !reused_bus {
        if let Some(events) = durable::load_events(&graph_id) {
            for event in events {
                bus.emit(event);
            }
        }
    }

    let state = Arc::new(existing_state.unwrap_or_default());
    let mut snapshot = snapshot;
    // Mark running before the async task is scheduled so resume returns a live status.
    graph::mark_running(&mut snapshot);
    let _ = durable::persist_graph(&snapshot);
    let snap_arc = Arc::new(RwLock::new(snapshot));
    let cancel = cancel::new_scope();

    let exec = Arc::new(GraphExecution {
        snapshot: snap_arc.clone(),
        bus: bus.clone(),
        state: state.clone(),
        worker: worker_name.clone(),
        input: input.clone(),
        profile: profile.clone(),
        plan: plan.clone(),
        cancel: cancel.clone(),
    });

    GRAPHS.write().insert(graph_id.0.clone(), exec);

    let record = ExecutionRecord {
        graph_id: graph_id.clone(),
        worker: worker_name.clone(),
        input: input.clone(),
        plan: plan.clone(),
        profile: profile.clone(),
        paused: false,
        cancelled: false,
    };
    let _ = durable::save_execution_record(&record);

    let gid = graph_id.clone();
    task::spawn(async move {
        run_worker(
            gid,
            worker_name,
            input,
            run,
            bus,
            state,
            snap_arc,
            profile,
            plan,
            cancel,
        )
        .await;
    });

    Ok(())
}

fn restore_exec_from_durable(id: &GraphId) -> Option<Arc<GraphExecution>> {
    let record = durable::load_execution_record(id)?;
    let snapshot = durable::load_graph(id)?;
    let state = durable::load_checkpoint(id).unwrap_or_default();
    let bus = Arc::new(EventBus::new());
    if let Some(events) = durable::load_events(id) {
        for event in events {
            bus.emit(event);
        }
    }
    Some(Arc::new(GraphExecution {
        snapshot: Arc::new(RwLock::new(snapshot)),
        bus,
        state: Arc::new(state),
        worker: record.worker,
        input: record.input,
        profile: record.profile,
        plan: record.plan,
        cancel: cancel::new_scope(),
    }))
}

#[allow(clippy::too_many_arguments)]
async fn run_worker(
    graph_id: GraphId,
    _name: String,
    input: Value,
    run: WorkerFn,
    bus: SharedEventBus,
    state: Arc<StateStore>,
    snapshot: Arc<RwLock<GraphSnapshot>>,
    profile: ResourceProfile,
    plan: super::types::ExecutionPlan,
    cancel: CancellationToken,
) {
    {
        let mut snap = snapshot.write();
        // Already marked running in spawn_execution; keep durable in sync if needed.
        if snap.status != GraphStatus::Running {
            graph::mark_running(&mut snap);
            let _ = durable::persist_graph(&snap);
        }
    }

    let ctx = WorkerContext::new(
        graph_id.clone(),
        bus.clone(),
        state.clone(),
        snapshot.clone(),
        cancel.clone(),
    );
    emit_worker_start(&ctx);
    ctx.log("execution started");

    let started = super::id::now_ms();
    let result = runner::run_with_plan(
        &plan,
        input,
        run,
        graph_id.clone(),
        bus.clone(),
        state.clone(),
        snapshot.clone(),
        profile,
        cancel.clone(),
    )
    .await;

    match result {
        Ok(value) => {
            bus.emit(emit::result(value.clone()));
            let duration = super::id::now_ms().saturating_sub(started);
            bus.emit(emit::node_done(ctx.node_id.clone(), duration));
            let mut snap = snapshot.write();
            graph::mark_done(&mut snap);
            let _ = durable::persist_graph(&snap);
            let _ = durable::persist_events(&graph_id, &bus.history());
            let final_snap = snap.clone();
            drop(snap);
            super::metrics::inc_graph_completed();
            super::webhooks::notify_done(&final_snap, duration, Some(value));
            bus.emit(emit::graph_done(graph_id.clone()));
            release_live_graph(&graph_id);
        }
        Err(ResumaError::Cancelled) => {
            let duration = super::id::now_ms().saturating_sub(started);
            let hard_cancel = durable::load_execution_record(&graph_id)
                .map(|r| r.cancelled)
                .unwrap_or(false);
            // Resume replaces GRAPHS with a fresh snapshot Arc — skip durable
            // writes so the soft-pause finalizer cannot clobber the new run.
            let superseded = GRAPHS
                .read()
                .get(&graph_id.0)
                .map(|e| !Arc::ptr_eq(&e.snapshot, &snapshot))
                .unwrap_or(true);
            if hard_cancel {
                if superseded {
                    return;
                }
                ctx.log("execution cancelled");
                let mut snap = snapshot.write();
                graph::mark_failed(&mut snap);
                let _ = durable::persist_graph(&snap);
                let _ = durable::persist_events(&graph_id, &bus.history());
                let final_snap = snap.clone();
                drop(snap);
                super::metrics::inc_graph_failed();
                super::webhooks::notify_failed(
                    &final_snap,
                    duration,
                    "cancelled by operator".into(),
                );
                bus.emit(emit::graph_done(graph_id.clone()));
                release_live_graph(&graph_id);
            } else if superseded {
                // A resume already took over — leave the new execution alone.
            } else {
                ctx.log("execution paused (cancelled)");
                let mut snap = snapshot.write();
                graph::mark_paused(&mut snap);
                let _ = durable::persist_graph(&snap);
                let _ = durable::persist_events(&graph_id, &bus.history());
                let final_snap = snap.clone();
                drop(snap);
                super::metrics::inc_graph_paused();
                super::webhooks::notify_paused(&final_snap, duration);
            }
        }
        Err(e) => {
            bus.emit(emit::node_failed(ctx.node_id.clone(), e.to_string()));
            let duration = super::id::now_ms().saturating_sub(started);
            let mut snap = snapshot.write();
            graph::mark_failed(&mut snap);
            let _ = durable::persist_graph(&snap);
            let _ = durable::persist_events(&graph_id, &bus.history());
            let final_snap = snap.clone();
            let err = e.to_string();
            drop(snap);
            super::metrics::inc_graph_failed();
            super::webhooks::notify_failed(&final_snap, duration, err);
            bus.emit(emit::graph_done(graph_id.clone()));
            release_live_graph(&graph_id);
        }
    }
}

/// Newest-first scan for `field` in the last `WorkerEvent::Result` payload.
/// Pulled out of [`FlowEngine::last_result_field`] so it is testable without a
/// live graph/worker registry.
fn extract_last_result_field(events: &[super::types::WorkerEvent], field: &str) -> Option<String> {
    events.iter().rev().find_map(|e| match e {
        super::types::WorkerEvent::Result { data, .. } => {
            data.get(field).and_then(json_field_as_string)
        }
        _ => None,
    })
}

/// Coerce a JSON string / number / bool to `String`. Objects, arrays, and null
/// return `None` — those need a typed deserialize path, not a field helper.
fn json_field_as_string(v: &serde_json::Value) -> Option<String> {
    match v {
        serde_json::Value::String(s) => Some(s.clone()),
        serde_json::Value::Number(n) => Some(n.to_string()),
        serde_json::Value::Bool(b) => Some(b.to_string()),
        _ => None,
    }
}

#[cfg(test)]
mod last_artifact_tests {
    use super::*;
    use crate::exec::types::{NodeId, WorkerEvent};
    use serde_json::json;

    #[test]
    fn finds_artifact_id_from_most_recent_result_only() {
        let events = vec![
            WorkerEvent::NodeStart {
                node: NodeId("n0".into()),
                kind: crate::exec::types::NodeKind::Worker,
                timestamp_ms: 0,
            },
            WorkerEvent::Result {
                data: json!({ "artifact_id": "art-old" }),
                timestamp_ms: 1,
            },
            WorkerEvent::Result {
                data: json!({ "artifact_id": "art-new", "bytes": 1024 }),
                timestamp_ms: 2,
            },
        ];
        assert_eq!(
            extract_last_result_field(&events, "artifact_id"),
            Some("art-new".to_string())
        );
    }

    #[test]
    fn returns_none_without_a_result_event_or_matching_field() {
        let events = vec![WorkerEvent::Log {
            message: "hi".into(),
            node: NodeId("n0".into()),
            timestamp_ms: 0,
        }];
        assert_eq!(extract_last_result_field(&events, "artifact_id"), None);

        let events = vec![WorkerEvent::Result {
            data: json!({ "chunk_id": "c1" }),
            timestamp_ms: 0,
        }];
        assert_eq!(extract_last_result_field(&events, "artifact_id"), None);
        assert_eq!(
            extract_last_result_field(&events, "chunk_id"),
            Some("c1".to_string())
        );
    }

    #[test]
    fn coerces_number_and_bool_fields_to_string() {
        let events = vec![WorkerEvent::Result {
            data: json!({ "bytes": 1024, "ok": true, "nested": { "x": 1 } }),
            timestamp_ms: 0,
        }];
        assert_eq!(
            extract_last_result_field(&events, "bytes"),
            Some("1024".to_string())
        );
        assert_eq!(
            extract_last_result_field(&events, "ok"),
            Some("true".to_string())
        );
        assert_eq!(extract_last_result_field(&events, "nested"), None);
    }
}