nornir 0.5.4

//! **Executor — drain the plan-DAG.**
//!
//! The *planner* ([`super::decompose`] / `nornir funnel prompt`) turns an idea
//! into a plan-DAG. The **executor** walks that DAG's *ready frontier* and does
//! the work: for each node whose dependencies are all `Done`, it drives the node
//! through the existing generate→judge loop ([`super::decompose::run_subtask_loop`]
//! — a [`Generator`] proposes, an [`AnswerJudge`] oracle accepts/rejects), records
//! the node `Done`|`Failed` ([`Event::NodeStatusChanged`]), re-promotes the
//! frontier, and repeats until the DAG **drains** (every node `Done`) or **blocks**
//! (a `Failed` node strands its dependents — surfaced, never a silent stall).
//!
//! **Planner and executor are SEPARATE roles.** The executor consumes the
//! `Generator` seam and never re-uses the planner's decomposition call; a
//! Claude-backed `Generator` (when it lands) plugs in here without touching the
//! planner's Claude call. Each node is one unit of work — the natural granularity
//! for a future **jera** job fanned across worker machines.

use anyhow::{anyhow, Result};
use chrono::Utc;
use std::collections::BTreeMap;

use crate::warehouse::codegen_judge::AnswerJudge;
use crate::warehouse::generator::{GenRequest, Generator};

use super::decompose::{run_subtask_loop, RetrievedContext, Subtask};
use super::event::{Event, NodeStatus};
use super::ids::{NodeId, PlanId};
use super::state::{Funnel, PlanNode};
use super::store::Store;

/// Why a drain stopped.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DrainOutcome {
    /// Every node in the plan reached `Done`.
    Drained,
    /// The frontier emptied with work left un-`Done` (a `Failed` node stranded
    /// its dependents, or the plan was already stuck).
    Blocked,
}

/// The result of draining one plan-DAG.
#[derive(Debug, Clone)]
pub struct DrainReport {
    pub plan_id: String,
    /// Nodes the executor actually ran.
    pub ran: usize,
    /// Nodes the oracle accepted (→ `Done`).
    pub done: usize,
    /// Nodes the oracle rejected (→ `Failed`).
    pub failed: usize,
    /// Node ids left un-`Done` when the drain stopped, id order.
    pub blocked: Vec<String>,
    pub outcome: DrainOutcome,
}

impl DrainReport {
    /// Human-readable summary (plain numbers — a nordisk law).
    pub fn render(&self) -> String {
        let head = match self.outcome {
            DrainOutcome::Drained => "drained (all nodes done)",
            DrainOutcome::Blocked => "BLOCKED (frontier empty with work left)",
        };
        let mut s = format!(
            "plan {} — {}\n  ran {} · done {} · failed {}\n",
            self.plan_id, head, self.ran, self.done, self.failed
        );
        if !self.blocked.is_empty() {
            s.push_str(&format!("  blocked/pending: {}\n", self.blocked.join(", ")));
        }
        s
    }
}

/// Reconstruct a runnable [`Subtask`] from a plan node (title / target / anchor).
/// The anchor is left empty here — the executor drives the node by its title +
/// parent task; RAG anchoring is the planner's job (see `anchor_subtasks`).
fn subtask_from_node(index: usize, node: &PlanNode) -> Subtask {
    let title = node
        .params
        .get("title")
        .and_then(|v| v.as_str())
        .map(str::to_string)
        .or_else(|| node.prompt_excerpt.clone())
        .unwrap_or_else(|| node.id.as_str().to_string());
    let target_file = node.targets.first().cloned().unwrap_or_default();
    Subtask {
        index,
        title: title.clone(),
        target_file,
        anchor_symbol: String::new(),
        acceptance: "compiler/test oracle accepts".into(),
        context: RetrievedContext { query: title, anchors: Vec::new(), retrievers: Vec::new() },
    }
}

/// **Drain `plan_id`'s DAG.** Repeatedly: promote the ready frontier, run every
/// ready node through the generate→judge loop, record `Done`|`Failed`, until the
/// frontier empties. Returns a [`DrainReport`]; `Blocked` when any node is left
/// un-`Done` (a failed node strands its dependents — never a silent stall).
///
/// Offline-capable: pass a mock/canned [`Generator`] + any [`AnswerJudge`]. Each
/// node run maps to a future jera job (fan-out across worker machines).
pub fn drain_plan(
    store: &mut Store,
    plan_id: &PlanId,
    generator: &dyn Generator,
    judge: &dyn AnswerJudge,
    max_rounds: usize,
) -> Result<DrainReport> {
    // The parent task text grounds each node's prompt (the mother idea, or the
    // plan summary as a fallback).
    let parent_task = {
        let plan = store
            .funnel
            .plans
            .get(plan_id)
            .ok_or_else(|| anyhow!("no plan {}", plan_id.as_str()))?;
        store
            .funnel
            .ideas
            .get(&plan.idea_id)
            .map(|i| i.text.clone())
            .unwrap_or_else(|| plan.summary.clone())
    };

    let (mut ran, mut done, mut failed) = (0usize, 0usize, 0usize);

    loop {
        store.funnel.promote_ready();
        // Snapshot this wave's ready nodes (owned) so the mutable record() calls
        // below don't overlap the immutable plan borrow.
        let wave: Vec<(NodeId, Subtask)> = {
            let plan = store
                .funnel
                .plans
                .get(plan_id)
                .ok_or_else(|| anyhow!("plan {} vanished mid-drain", plan_id.as_str()))?;
            plan.nodes
                .values()
                .filter(|n| n.status == NodeStatus::Ready)
                .enumerate()
                .map(|(i, n)| (n.id.clone(), subtask_from_node(i, n)))
                .collect()
        };
        if wave.is_empty() {
            break;
        }
        for (node_id, sub) in wave {
            store.record(Event::NodeStatusChanged {
                plan_id: plan_id.clone(),
                node_id: node_id.clone(),
                status: NodeStatus::InProgress,
                why: Some("executor: running".into()),
                ts: Utc::now(),
            })?;
            let verdict = run_subtask_loop(&parent_task, &sub, generator, judge, max_rounds)?;
            ran += 1;
            let status = verdict.node_status();
            if status == NodeStatus::Done {
                done += 1;
            } else {
                failed += 1;
            }
            store.record(Event::NodeStatusChanged {
                plan_id: plan_id.clone(),
                node_id,
                status,
                why: Some(format!(
                    "executor: {} (score {:.2}, {}/{} tests)",
                    verdict.best.message,
                    verdict.best.score,
                    verdict.best.tests_passed,
                    verdict.best.tests_total
                )),
                ts: Utc::now(),
            })?;
        }
    }

    let plan = store
        .funnel
        .plans
        .get(plan_id)
        .ok_or_else(|| anyhow!("plan {} vanished", plan_id.as_str()))?;
    let blocked: Vec<String> = plan
        .nodes
        .values()
        .filter(|n| n.status != NodeStatus::Done)
        .map(|n| n.id.as_str().to_string())
        .collect();
    let outcome = if blocked.is_empty() { DrainOutcome::Drained } else { DrainOutcome::Blocked };
    Ok(DrainReport {
        plan_id: plan_id.as_str().to_string(),
        ran,
        done,
        failed,
        blocked,
        outcome,
    })
}

/// **Dry-run:** the topological *waves* the executor would run for `plan_id`,
/// without executing anything (nodes already `Failed` and their dependents are
/// omitted). Pure — drives the frontier logic for the CLI/UI without a generator
/// or oracle.
pub fn plan_frontier_waves(funnel: &Funnel, plan_id: &PlanId) -> Vec<Vec<String>> {
    let Some(plan) = funnel.plans.get(plan_id) else {
        return Vec::new();
    };
    let mut status: BTreeMap<NodeId, NodeStatus> =
        plan.nodes.iter().map(|(id, n)| (id.clone(), n.status)).collect();
    let mut waves = Vec::new();
    loop {
        let mut wave: Vec<NodeId> = Vec::new();
        for (id, st) in &status {
            if matches!(st, NodeStatus::Done | NodeStatus::Failed) {
                continue;
            }
            let deps_done = plan
                .edges
                .iter()
                .filter(|(_, to)| to == id)
                .all(|(from, _)| matches!(status.get(from), Some(NodeStatus::Done)));
            if deps_done {
                wave.push(id.clone());
            }
        }
        if wave.is_empty() {
            break;
        }
        for id in &wave {
            status.insert(id.clone(), NodeStatus::Done);
        }
        waves.push(wave.iter().map(|n| n.as_str().to_string()).collect());
    }
    waves
}

#[cfg(test)]
mod tests {
    use super::{drain_plan, plan_frontier_waves, DrainOutcome};
    use super::super::decompose::{store_subtasks, Decomposition, RetrievedContext, Subtask};
    use super::super::event::{Event, ItemKind, NodeStatus};
    use super::super::ids::{IdeaId, PlanId};
    use super::super::store::Store;
    use crate::warehouse::codegen_judge::{AnswerJudge, Verdict};
    use crate::warehouse::generator::{Backend, GenAnswer, GenRequest, Generator};
    use anyhow::Result;
    use chrono::Utc;

    /// A generator that echoes the prompt — so the judge can decide per node (the
    /// node's title rides into the prompt via `build_codegen_prompt`).
    struct EchoGen;
    impl Backend for EchoGen {
        fn id(&self) -> &str {
            "echo"
        }
        fn available(&self) -> bool {
            true
        }
    }
    impl Generator for EchoGen {
        fn complete(&self, req: &GenRequest) -> Result<GenAnswer> {
            Ok(GenAnswer {
                text: req.prompt.clone(),
                tokens_in: 0,
                tokens_out: 0,
                tokens_per_s: 0.0,
                latency_ms: 0.0,
            })
        }
    }

    /// Oracle: reject iff the answer carries the `FAILNODE` marker; else accept.
    struct MarkerJudge;
    impl AnswerJudge for MarkerJudge {
        fn judge(&self, answer: &str) -> Result<Verdict> {
            if answer.contains("FAILNODE") {
                Ok(Verdict::rejected("marked to fail"))
            } else {
                Ok(Verdict { compiled: true, tests_passed: 1, tests_total: 1, score: 1.0, message: "ok".into() })
            }
        }
    }

    fn sub(index: usize, title: &str) -> Subtask {
        Subtask {
            index,
            title: title.into(),
            target_file: format!("src/n{index}.rs"),
            anchor_symbol: String::new(),
            acceptance: "green".into(),
            context: RetrievedContext { query: title.into(), anchors: vec![], retrievers: vec![] },
        }
    }

    /// In-memory funnel with a chained plan of the given subtask titles.
    fn plan_with(titles: &[&str]) -> (Store, PlanId) {
        let dir = tempfile::tempdir().unwrap().into_path();
        let mut store = Store::open(dir.join("wh")).unwrap();
        let idea = IdeaId::seq(store.funnel.next_idea.max(1));
        store
            .record(Event::IdeaSubmitted {
                id: idea.clone(),
                source: "test".into(),
                text: "parent task".into(),
                refs: vec![],
                item_kind: ItemKind::Idea,
                ts: Utc::now(),
            })
            .unwrap();
        let decomp = Decomposition {
            task: "parent task".into(),
            subtasks: titles.iter().enumerate().map(|(i, t)| sub(i, t)).collect(),
        };
        let (plan_id, _nodes) = store_subtasks(&mut store, &idea, &decomp).unwrap();
        (store, plan_id)
    }

    /// Headless drive of the WHOLE funnel flow (the coverage LAW): a planner built
    /// a 3-node plan-DAG; the executor drains it to all-Done offline (mock gen +
    /// accepting oracle). Asserts the plan-DAG grew AND drained.
    #[test]
    fn drain_reaches_all_done_offline() {
        let (mut store, plan_id) = plan_with(&["node a", "node b", "node c"]);
        assert_eq!(store.funnel.plans.get(&plan_id).unwrap().nodes.len(), 3, "plan-DAG grew to 3 nodes");
        let rep = drain_plan(&mut store, &plan_id, &EchoGen, &MarkerJudge, 1).unwrap();
        assert_eq!(rep.outcome, DrainOutcome::Drained);
        assert_eq!((rep.ran, rep.done, rep.failed), (3, 3, 0));
        assert!(rep.blocked.is_empty());
        assert!(store
            .funnel
            .plans
            .get(&plan_id)
            .unwrap()
            .nodes
            .values()
            .all(|n| n.status == NodeStatus::Done));
    }

    /// A failed node strands its dependent → the drain BLOCKS (never silent stall).
    #[test]
    fn drain_blocks_when_a_node_fails_and_strands_dependents() {
        let (mut store, plan_id) = plan_with(&["node a", "FAILNODE b", "node c"]);
        let rep = drain_plan(&mut store, &plan_id, &EchoGen, &MarkerJudge, 1).unwrap();
        assert_eq!(rep.outcome, DrainOutcome::Blocked);
        assert_eq!(rep.done, 1, "only the root ran green");
        assert_eq!(rep.failed, 1, "the marked node failed");
        assert_eq!(rep.blocked.len(), 2, "the failed node + its stranded dependent");
    }

    /// The dry-run frontier walk is topological (chained plan → one node/wave).
    #[test]
    fn frontier_waves_are_topological() {
        let (store, plan_id) = plan_with(&["a", "b", "c"]);
        let waves = plan_frontier_waves(&store.funnel, &plan_id);
        assert_eq!(waves.len(), 3);
        assert!(waves.iter().all(|w| w.len() == 1));
    }
}