nornir 0.4.28

//! **Funnel task decomposition with the compiler-as-oracle** (EPIC #45).
//!
//! A funnel *task* (an idea / error / test item) is turned into ~N atomic,
//! testable **subtasks**, each anchored to a concrete file/fn + an acceptance
//! check, then each subtask is driven through a bounded *generate → judge →
//! retry* loop where the [`CodegenJudge`](crate::warehouse::codegen_judge) is the
//! accept/reject **oracle** (apply the patch, `cargo build` + test, score). No
//! cloud, no big model: a local [`Generator`](crate::warehouse::generator) (the
//! `mock` backend for tests/offline, a real `--gen <spec>` backend otherwise)
//! proposes the code; the compiler decides if it's right.
//!
//! # The pipeline
//!
//! ```text
//!  task text ─┬─ retrieve_context (RAG) ──→ anchors (file/fn snippets)
//!             │      ├ code index (symbol_facts/call_edges)  ← always
//!             │      └ vector index (Embedder)               ← #[cfg(vector)]
//!             ├─ decompose ──────────────→ Vec<Subtask>  (≤20, each anchored)
//!             ├─ store_subtasks ─────────→ child funnel DAG nodes under parent
//!             └─ for each subtask:
//!                  run_subtask_loop ─ Generator → CodegenJudge oracle
//!                    accept (judge 1.0) → node Done
//!                    reject (judge 0.0/partial) → retry w/ compiler error
//!                                                  + more context, bounded
//!                    exhausted → node Failed (flagged)
//! ```
//!
//! This mirrors [`planner`](super::planner)'s spirit (RAG-grounded proposal via
//! the embedder + dep graph) but goes one step further: it proposes *atomic
//! subtasks* and *code patches*, and lets the compiler/tests be the judge.
//!
//! # The simulation
//!
//! [`simulate`] runs the WHOLE loop offline: `mock` Generator + a REAL
//! [`CodegenJudge`] against a tiny fixture crate, printing the subtask DAG +
//! per-subtask verdicts. It needs no model and no network — the accept/reject is
//! the compiler. The CLI exposes it as `nornir funnel decompose <task>`.
//!
//! # LAW 6 (readable state)
//!
//! Every produced value emits a `state_json` ([`Subtask::to_json`],
//! [`RetrievedContext::to_json`], [`SubtaskVerdict::to_json`],
//! [`Decomposition::to_json`], [`SimulationReport::to_json`]) so the viz/CLI can
//! "see what the user sees" as data, not pixels.

use anyhow::{anyhow, Result};
use chrono::Utc;
use serde::{Deserialize, Serialize};

use crate::knowledge::query::KnowledgeView;
use crate::knowledge::symbols::SymbolRow;
use crate::warehouse::codegen_judge::{AnswerJudge, Verdict};
use crate::warehouse::generator::{GenRequest, Generator};

use super::event::{Event, ItemKind, NodeStatus, PlanStatus};
use super::ids::{IdeaId, NodeId, PlanId};
use super::store::Store;

/// The hard cap on subtasks a single task may decompose into. The funnel is a
/// DAG and a runaway decomposition would blow it up — #45 fixes ≤20.
pub const MAX_SUBTASKS: usize = 20;

/// Default bounded retry rounds for a subtask's generate→judge loop (the first
/// attempt plus this many retries that re-feed the compiler error).
pub const DEFAULT_MAX_ROUNDS: usize = 3;

/// Default top-k anchors RAG retrieves per (sub)task.
pub const DEFAULT_TOP_K: usize = 5;

// ─── retrieval (RAG) ─────────────────────────────────────────────────────────

/// One retrieved code anchor: a concrete symbol (file/fn) the (sub)task is
/// grounded in, with an optional source excerpt and the retrieval score.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Anchor {
    /// The symbol/file the anchor points at (`item_name` for a code-index hit).
    pub symbol: String,
    /// Repo-relative file path.
    pub file: String,
    /// 1-based line of the symbol/snippet.
    pub line: u32,
    /// The item kind (`fn` / `struct` / …) for a code-index hit, or `chunk` for
    /// a vector-index hit.
    pub kind: String,
    /// Optional source excerpt (the symbol signature, or a vector chunk).
    pub excerpt: Option<String>,
    /// Retrieval score (relevance). Higher = more relevant.
    pub score: f32,
    /// Which retriever produced this anchor (`code_index` | `vector`).
    pub source: String,
}

impl Anchor {
    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "symbol": self.symbol,
            "file": self.file,
            "line": self.line,
            "kind": self.kind,
            "excerpt": self.excerpt,
            "score": self.score,
            "source": self.source,
        })
    }
}

/// The result of RAG retrieval for a (sub)task: the top-k anchors plus which
/// retrievers actually contributed (so a `--gen`/`vector`-less build is honest
/// about degrading to the code index).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RetrievedContext {
    /// The query text the anchors were retrieved for.
    pub query: String,
    /// Top-k anchors, best-first.
    pub anchors: Vec<Anchor>,
    /// Retrievers that contributed at least one anchor (`code_index`,
    /// `vector`).
    pub retrievers: Vec<String>,
}

impl RetrievedContext {
    /// Render the anchors as a compact prompt block the Generator can read.
    pub fn as_prompt_block(&self) -> String {
        if self.anchors.is_empty() {
            return "(no anchors retrieved)".to_string();
        }
        let mut s = String::new();
        for a in &self.anchors {
            s.push_str(&format!("- {} ({}:{}) [{}]", a.symbol, a.file, a.line, a.kind));
            if let Some(ex) = &a.excerpt {
                s.push_str(&format!("  — {}", ex.trim()));
            }
            s.push('\n');
        }
        s
    }

    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "query": self.query,
            "anchors": self.anchors.iter().map(Anchor::to_json).collect::<Vec<_>>(),
            "retrievers": self.retrievers,
        })
    }
}

/// The seam a retriever implements. The code index (always available) and the
/// vector index (`#[cfg(feature = "vector")]`) both implement it, so [`retrieve`]
/// can fan out over whatever is present and still produce a sane result with
/// neither a model nor a network.
pub trait Retriever {
    /// A short id for the `retrievers`/`source` fields.
    fn id(&self) -> &str;
    /// Return up to `k` anchors most relevant to `query`, best-first.
    fn retrieve(&self, query: &str, k: usize) -> Vec<Anchor>;
}

/// The **code-index retriever** over a [`KnowledgeView`] (`symbol_facts` +
/// `call_edges`). 100% offline + model-free: it scores a symbol by lexical token
/// overlap between the query and the symbol name/file/signature, so a query like
/// "retry the http fetch" anchors onto a `fn fetch_retry` in `http.rs`. This is
/// the always-present RAG leg (the vector leg layers on top when compiled).
pub struct CodeIndexRetriever<'a> {
    view: &'a KnowledgeView,
}

impl<'a> CodeIndexRetriever<'a> {
    pub fn new(view: &'a KnowledgeView) -> Self {
        Self { view }
    }
}

/// Lower-case alphanumeric tokens of `s` (split on any non-alnum), length ≥ 2 so
/// noise like single letters / separators doesn't drive the match.
fn tokens(s: &str) -> Vec<String> {
    s.split(|c: char| !c.is_alphanumeric())
        .filter(|t| t.len() >= 2)
        .map(|t| t.to_ascii_lowercase())
        .collect()
}

/// Lexical relevance of a symbol to a query's token set: the count of query
/// tokens that appear (as a substring of, or equal to) any of the symbol's own
/// tokens, weighted so a name hit outranks a file/signature hit.
fn lexical_score(query_tokens: &[String], sym: &SymbolRow) -> f32 {
    let name_tokens = tokens(&sym.item_name);
    let file_tokens = tokens(&sym.file);
    let sig_tokens = sym.signature.as_deref().map(tokens).unwrap_or_default();
    let mut score = 0.0f32;
    for qt in query_tokens {
        if name_tokens.iter().any(|t| t == qt || t.contains(qt.as_str())) {
            score += 3.0;
        } else if sig_tokens.iter().any(|t| t == qt) {
            score += 1.5;
        } else if file_tokens.iter().any(|t| t == qt || t.contains(qt.as_str())) {
            score += 1.0;
        }
    }
    score
}

impl Retriever for CodeIndexRetriever<'_> {
    fn id(&self) -> &str {
        "code_index"
    }
    fn retrieve(&self, query: &str, k: usize) -> Vec<Anchor> {
        let qt = tokens(query);
        let mut scored: Vec<(f32, &SymbolRow)> = self
            .view
            .symbols
            .iter()
            .map(|s| (lexical_score(&qt, s), s))
            .filter(|(score, _)| *score > 0.0)
            .collect();
        // Best-first; ties broken by (file, line) for determinism.
        scored.sort_by(|a, b| {
            b.0.partial_cmp(&a.0)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.1.file.cmp(&b.1.file))
                .then_with(|| a.1.line.cmp(&b.1.line))
        });
        scored
            .into_iter()
            .take(k)
            .map(|(score, s)| Anchor {
                symbol: s.item_name.clone(),
                file: s.file.clone(),
                line: s.line,
                kind: s.item_kind.clone(),
                excerpt: s.signature.clone(),
                score,
                source: "code_index".to_string(),
            })
            .collect()
    }
}

/// Fan out over the supplied retrievers, merge their anchors, dedup by
/// `(file, symbol)` keeping the highest score, and return the global top-k. Pure
/// over the retriever seam, so a test injects a fake retriever and asserts the
/// merge/dedup/rank without a warehouse.
pub fn retrieve(retrievers: &[&dyn Retriever], query: &str, k: usize) -> RetrievedContext {
    let mut all: Vec<Anchor> = Vec::new();
    let mut contributed: Vec<String> = Vec::new();
    for r in retrievers {
        let hits = r.retrieve(query, k);
        if !hits.is_empty() && !contributed.iter().any(|c| c == r.id()) {
            contributed.push(r.id().to_string());
        }
        all.extend(hits);
    }
    // Dedup by (file, symbol), keeping the best score.
    all.sort_by(|a, b| {
        a.file
            .cmp(&b.file)
            .then_with(|| a.symbol.cmp(&b.symbol))
            .then_with(|| {
                b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)
            })
    });
    all.dedup_by(|a, b| a.file == b.file && a.symbol == b.symbol);
    // Global rank by score (desc), tie-break (file, line) for determinism.
    all.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.file.cmp(&b.file))
            .then_with(|| a.line.cmp(&b.line))
    });
    all.truncate(k);
    RetrievedContext { query: query.to_string(), anchors: all, retrievers: contributed }
}

// ─── decomposition ───────────────────────────────────────────────────────────

/// One atomic, testable subtask of a decomposed task: a title, the anchor it is
/// grounded in (a concrete file/fn from RAG), an acceptance check, and a
/// retrieved-context handle (so the per-subtask loop can ground the generator).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Subtask {
    /// 0-based index within the decomposition (drives the child node order).
    pub index: usize,
    /// Imperative one-liner ("implement add() in src/lib.rs").
    pub title: String,
    /// The concrete file the subtask targets (from its anchor).
    pub target_file: String,
    /// The symbol the subtask is anchored to (the fn/struct it touches).
    pub anchor_symbol: String,
    /// The acceptance check that must pass for the subtask to be accepted (a
    /// human-readable description of what the [`CodegenJudge`] verifies).
    pub acceptance: String,
    /// The RAG context this subtask was decomposed against.
    pub context: RetrievedContext,
}

impl Subtask {
    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "index": self.index,
            "title": self.title,
            "target_file": self.target_file,
            "anchor_symbol": self.anchor_symbol,
            "acceptance": self.acceptance,
            "context": self.context.to_json(),
        })
    }
}

/// The full decomposition of one task: the parent text + the ordered subtasks.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Decomposition {
    /// The original task text.
    pub task: String,
    /// The atomic subtasks (≤ [`MAX_SUBTASKS`]), each anchored.
    pub subtasks: Vec<Subtask>,
}

impl Decomposition {
    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "task": self.task,
            "subtask_count": self.subtasks.len(),
            "subtasks": self.subtasks.iter().map(Subtask::to_json).collect::<Vec<_>>(),
        })
    }
}

/// Decompose `task` into ≤`max` atomic subtasks, each anchored to a concrete
/// file/fn retrieved from the supplied `retrievers` (RAG). One subtask is
/// produced per retrieved anchor (anchors ARE the atoms — each concrete symbol a
/// task touches is one testable unit), capped at `max`. With no anchors (empty
/// index) a single whole-task subtask is produced so the loop still runs.
///
/// This is deliberately model-free: the decomposition is *grounded in the code
/// index*, not hallucinated by an LLM — exactly the planner's spirit (the
/// dep-graph/index does the structural work; a model only fills code in later).
pub fn decompose(
    task: &str,
    retrievers: &[&dyn Retriever],
    top_k: usize,
    max: usize,
) -> Decomposition {
    let max = max.min(MAX_SUBTASKS).max(1);
    let ctx = retrieve(retrievers, task, top_k.max(max));
    let mut subtasks = Vec::new();
    for (i, anchor) in ctx.anchors.iter().take(max).enumerate() {
        // Per-subtask context: just this anchor (so the generator is grounded on
        // the one symbol it must touch), preserving the parent's retrievers.
        let sub_ctx = RetrievedContext {
            query: format!("{task} :: {}", anchor.symbol),
            anchors: vec![anchor.clone()],
            retrievers: ctx.retrievers.clone(),
        };
        subtasks.push(Subtask {
            index: i,
            title: format!("implement `{}` in {}", anchor.symbol, anchor.file),
            target_file: anchor.file.clone(),
            anchor_symbol: anchor.symbol.clone(),
            acceptance: format!(
                "cargo build + acceptance tests green after editing {} ({})",
                anchor.file, anchor.symbol
            ),
            context: sub_ctx,
        });
    }
    // Fallback: an empty index still yields one runnable subtask.
    if subtasks.is_empty() {
        subtasks.push(Subtask {
            index: 0,
            title: format!("implement task: {}", first_line(task)),
            target_file: String::new(),
            anchor_symbol: String::new(),
            acceptance: "cargo build + acceptance tests green".to_string(),
            context: ctx.clone(),
        });
    }
    Decomposition { task: task.to_string(), subtasks }
}

fn first_line(s: &str) -> String {
    s.lines().next().unwrap_or("").trim().to_string()
}

// ─── the per-subtask generate → judge → retry loop ──────────────────────────

/// One round of a subtask's loop: the prompt sent, the answer produced, and the
/// oracle's [`Verdict`]. Kept so the report explains *why* a subtask passed/
/// failed (which round, what the compiler said).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Round {
    /// 1-based round number.
    pub round: usize,
    /// First line of the prompt (for the trace — full prompt is large).
    pub prompt_excerpt: String,
    /// The judge's verdict for this round's answer.
    pub verdict: Verdict,
}

impl Round {
    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "round": self.round,
            "prompt_excerpt": self.prompt_excerpt,
            "verdict": self.verdict,
        })
    }
}

/// The outcome of a subtask's whole bounded loop: accepted/failed, the winning
/// verdict, and every round's verdict (the audit trail).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SubtaskVerdict {
    /// The subtask index this verdict is for.
    pub subtask_index: usize,
    /// Did the oracle accept any round (compiled + all acceptance tests pass)?
    pub accepted: bool,
    /// The best verdict seen (the accepting one, or the highest-scoring reject).
    pub best: Verdict,
    /// Every round's prompt/verdict, in order.
    pub rounds: Vec<Round>,
}

impl SubtaskVerdict {
    /// The funnel node status this verdict maps to: accepted → Done, else Failed.
    pub fn node_status(&self) -> NodeStatus {
        if self.accepted {
            NodeStatus::Done
        } else {
            NodeStatus::Failed
        }
    }

    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "subtask_index": self.subtask_index,
            "accepted": self.accepted,
            "best": self.best,
            "rounds": self.rounds.iter().map(Round::to_json).collect::<Vec<_>>(),
        })
    }
}

/// Build the code-generation prompt for a subtask round. On the first round it
/// carries the task + the subtask's anchor context; on a retry it ALSO carries
/// the previous round's compiler/test error so the generator can fix it (the
/// "compiler-as-oracle feedback" loop).
pub fn build_codegen_prompt(
    parent_task: &str,
    sub: &Subtask,
    prev_error: Option<&str>,
) -> String {
    let mut p = String::new();
    p.push_str("You are writing Rust code for one atomic subtask. Reply with ONLY ");
    p.push_str("the new file body (no prose, no fences).\n\n");
    p.push_str(&format!("Parent task: {parent_task}\n"));
    p.push_str(&format!("Subtask: {}\n", sub.title));
    p.push_str(&format!("Target file: {}\n", sub.target_file));
    p.push_str(&format!("Acceptance: {}\n\n", sub.acceptance));
    p.push_str("Relevant code (RAG anchors):\n");
    p.push_str(&sub.context.as_prompt_block());
    if let Some(err) = prev_error {
        p.push_str("\nThe previous attempt FAILED the compiler/tests:\n");
        p.push_str(err.trim());
        p.push_str("\nFix it.\n");
    }
    p
}

/// Drive ONE subtask through the bounded generate → judge → retry loop.
///
/// Each round: build the prompt (re-feeding the previous compiler error on a
/// retry), ask the [`Generator`] for a patch/file body, evaluate it with the
/// [`AnswerJudge`] oracle. Accept (and stop) on the first verdict that
/// [`Verdict::accepted`]; otherwise re-feed the verdict's `message` as the error
/// and retry, up to `max_rounds`. Returns the [`SubtaskVerdict`] (accepted/
/// failed + the full round trail).
///
/// The judge is the ORACLE — a non-compiling answer is a `0.0` verdict (a
/// reject), never a hard error; an `Err` here means the generator itself failed
/// (model missing, etc.).
pub fn run_subtask_loop(
    parent_task: &str,
    sub: &Subtask,
    generator: &dyn Generator,
    judge: &dyn AnswerJudge,
    max_rounds: usize,
) -> Result<SubtaskVerdict> {
    let max_rounds = max_rounds.max(1);
    let mut rounds = Vec::new();
    let mut prev_error: Option<String> = None;
    let mut best: Option<Verdict> = None;

    for round in 1..=max_rounds {
        let prompt = build_codegen_prompt(parent_task, sub, prev_error.as_deref());
        let req = GenRequest::new(prompt.clone()).with_max_tokens(1024);
        let answer = generator
            .complete(&req)
            .map_err(|e| anyhow!("generator failed on subtask {}: {e}", sub.index))?;
        let verdict = judge.judge(&answer.text)?;

        rounds.push(Round {
            round,
            prompt_excerpt: first_line(&prompt),
            verdict: verdict.clone(),
        });

        let improved = best
            .as_ref()
            .map(|b| verdict.score > b.score)
            .unwrap_or(true);
        if improved {
            best = Some(verdict.clone());
        }

        if verdict.accepted() {
            return Ok(SubtaskVerdict {
                subtask_index: sub.index,
                accepted: true,
                best: verdict,
                rounds,
            });
        }
        // Re-feed the compiler/test error for the next round.
        prev_error = Some(verdict.message.clone());
    }

    Ok(SubtaskVerdict {
        subtask_index: sub.index,
        accepted: false,
        best: best.unwrap_or_else(|| Verdict::rejected("no rounds run")),
        rounds,
    })
}

// ─── persist subtasks as child funnel DAG nodes ─────────────────────────────

/// Weave a [`Decomposition`] into the funnel as a child plan under `idea_id`:
/// one [`super::Plan`] whose nodes are the subtasks (in order, chained by
/// dependency edges so they dispatch top-down), each `targets = [target_file]`,
/// `node_kind` from the parent item kind. Returns the new [`PlanId`] + the node
/// ids in subtask order.
///
/// This is the structural half (no generation): the loop's verdicts are written
/// back with [`apply_verdicts`] once the subtasks have run.
pub fn store_subtasks(
    store: &mut Store,
    idea_id: &IdeaId,
    decomp: &Decomposition,
) -> Result<(PlanId, Vec<NodeId>)> {
    let kind = store
        .funnel
        .ideas
        .get(idea_id)
        .map(|i| i.item_kind)
        .unwrap_or(ItemKind::Idea);
    let node_kind = kind.node_kind().to_string();

    // Microsecond-stepped clock so replay order is the emit order.
    let base = Utc::now();
    let mut tick: i64 = 0;
    let mut at = || {
        let t = base + chrono::Duration::microseconds(tick);
        tick += 1;
        t
    };

    let plan_id = PlanId::seq(store.funnel.next_plan.max(1));
    store.record(Event::PlanCreated {
        id: plan_id.clone(),
        idea_id: idea_id.clone(),
        summary: format!(
            "decompose ({}) → {} atomic subtask(s)",
            kind.as_str(),
            decomp.subtasks.len()
        ),
        planner: "decompose".into(),
        ts: at(),
    })?;
    store.record(Event::PlanStatusChanged {
        plan_id: plan_id.clone(),
        status: PlanStatus::Active,
        why: Some("auto-planned from task decomposition".into()),
        ts: at(),
    })?;

    let mut node_ids = Vec::with_capacity(decomp.subtasks.len());
    let mut prev: Option<NodeId> = None;
    for sub in &decomp.subtasks {
        let node_id = NodeId::seq(store.funnel.next_node.max(1));
        let mut params = serde_json::Map::new();
        params.insert("title".into(), serde_json::Value::String(sub.title.clone()));
        params.insert("subtask_index".into(), serde_json::Value::from(sub.index));
        params.insert("acceptance".into(), serde_json::Value::String(sub.acceptance.clone()));
        params.insert("anchor".into(), serde_json::Value::String(sub.anchor_symbol.clone()));
        params.insert("state_json".into(), sub.to_json());
        let targets = if sub.target_file.is_empty() {
            vec![]
        } else {
            vec![sub.target_file.clone()]
        };
        store.record(Event::NodeAdded {
            plan_id: plan_id.clone(),
            node_id: node_id.clone(),
            kind: node_kind.clone(),
            params,
            targets,
            prompt_excerpt: Some(sub.title.clone()),
            ts: at(),
        })?;
        if let Some(from) = prev {
            store.record(Event::EdgeAdded {
                plan_id: plan_id.clone(),
                from_node: from,
                to_node: node_id.clone(),
                ts: at(),
            })?;
        }
        prev = Some(node_id.clone());
        node_ids.push(node_id);
    }
    store.funnel.promote_ready();
    Ok((plan_id, node_ids))
}

/// Write each subtask's loop verdict back onto its child node as a
/// [`Event::NodeStatusChanged`] (Done on accept, Failed on reject — a "flagged"
/// failure). `node_ids[i]` is the node for the subtask whose `subtask_index ==
/// i`. Returns the count written.
pub fn apply_verdicts(
    store: &mut Store,
    plan_id: &PlanId,
    node_ids: &[NodeId],
    verdicts: &[SubtaskVerdict],
) -> Result<usize> {
    let base = Utc::now();
    let mut tick = 0i64;
    let mut at = || {
        let t = base + chrono::Duration::microseconds(tick);
        tick += 1;
        t
    };
    let mut n = 0;
    for v in verdicts {
        let node_id = node_ids
            .get(v.subtask_index)
            .ok_or_else(|| anyhow!("verdict for subtask {} has no node", v.subtask_index))?;
        store.record(Event::NodeStatusChanged {
            plan_id: plan_id.clone(),
            node_id: node_id.clone(),
            status: v.node_status(),
            why: Some(format!(
                "judge: {} (score {:.2}, {}/{} tests)",
                v.best.message, v.best.score, v.best.tests_passed, v.best.tests_total
            )),
            ts: at(),
        })?;
        n += 1;
    }
    Ok(n)
}

// ─── the offline simulation ──────────────────────────────────────────────────

/// The end-to-end report of a [`simulate`] run: the decomposition + every
/// subtask's verdict + the aggregate. Everything needed to print the subtask DAG
/// with per-subtask verdicts, and to drive a `test_results` matrix row.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SimulationReport {
    /// The decomposition that ran.
    pub decomposition: Decomposition,
    /// One verdict per subtask, in subtask order.
    pub verdicts: Vec<SubtaskVerdict>,
    /// Subtasks accepted by the oracle.
    pub accepted: usize,
    /// Total subtasks.
    pub total: usize,
}

impl SimulationReport {
    /// Did every subtask pass the oracle?
    pub fn all_accepted(&self) -> bool {
        self.total > 0 && self.accepted == self.total
    }

    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "decomposition": self.decomposition.to_json(),
            "verdicts": self.verdicts.iter().map(SubtaskVerdict::to_json).collect::<Vec<_>>(),
            "accepted": self.accepted,
            "total": self.total,
        })
    }

    /// Human-readable subtask DAG + per-subtask verdicts (the CLI output).
    pub fn render(&self) -> String {
        let mut s = String::new();
        s.push_str(&format!(
            "task: {}\nsubtasks: {} ({} accepted by the compiler-oracle)\n",
            first_line(&self.decomposition.task),
            self.total,
            self.accepted
        ));
        for (sub, v) in self.decomposition.subtasks.iter().zip(&self.verdicts) {
            let mark = if v.accepted { "PASS" } else { "FAIL" };
            let dep = if sub.index == 0 {
                "(root)".to_string()
            } else {
                format!("(needs #{})", sub.index - 1)
            };
            s.push_str(&format!(
                "  #{} {} {}  {}\n      target: {} | anchor: {}\n      verdict: score {:.2} — {} (rounds: {})\n",
                sub.index,
                mark,
                sub.title,
                dep,
                if sub.target_file.is_empty() { "-" } else { &sub.target_file },
                if sub.anchor_symbol.is_empty() { "-" } else { &sub.anchor_symbol },
                v.best.score,
                v.best.message,
                v.rounds.len(),
            ));
        }
        s
    }
}

/// Run the WHOLE decomposition loop offline against an in-memory
/// [`KnowledgeView`] (the RAG anchor) + a [`Generator`] (use `mock` for offline)
/// + a REAL [`AnswerJudge`] oracle. Decomposes `task`, then runs each subtask's
/// bounded generate→judge→retry loop, and returns the [`SimulationReport`].
///
/// No model and no network are required: pass `generator("mock")?` and a
/// [`CodegenJudge`](crate::warehouse::codegen_judge::CodegenJudge) over a fixture
/// crate — the compiler is the oracle.
pub fn simulate(
    task: &str,
    view: &KnowledgeView,
    generator: &dyn Generator,
    judge: &dyn AnswerJudge,
    top_k: usize,
    max_subtasks: usize,
    max_rounds: usize,
) -> Result<SimulationReport> {
    let code = CodeIndexRetriever::new(view);
    let retrievers: Vec<&dyn Retriever> = vec![&code];
    let decomp = decompose(task, &retrievers, top_k, max_subtasks);

    let mut verdicts = Vec::with_capacity(decomp.subtasks.len());
    let mut accepted = 0;
    for sub in &decomp.subtasks {
        let v = run_subtask_loop(task, sub, generator, judge, max_rounds)?;
        if v.accepted {
            accepted += 1;
        }
        verdicts.push(v);
    }
    let total = decomp.subtasks.len();
    Ok(SimulationReport { decomposition: decomp, verdicts, accepted, total })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::knowledge::symbols::{CallEdgeRow, SymbolRow};
    use crate::warehouse::codegen_judge::{CodegenJudge, CodegenSpec};
    use crate::warehouse::generator::{Backend, GenAnswer, MockGenerator};

    // ── fixtures ────────────────────────────────────────────────────────

    fn sym(name: &str, kind: &str, file: &str, line: u32, sig: &str) -> SymbolRow {
        SymbolRow {
            crate_name: "fixture".into(),
            module_path: "fixture".into(),
            item_kind: kind.into(),
            item_name: name.into(),
            visibility: "pub".into(),
            file: file.into(),
            line,
            doc_lines: 0,
            signature: Some(sig.into()),
        }
    }

    /// A tiny code index: an `add` fn, a `retry_fetch` fn, a `Config` struct.
    fn view() -> KnowledgeView {
        KnowledgeView {
            symbols: vec![
                sym("add", "fn", "src/math.rs", 3, "fn add(a: i32, b: i32) -> i32"),
                sym("retry_fetch", "fn", "src/http.rs", 10, "fn retry_fetch(url: &str) -> Result<String>"),
                sym("Config", "struct", "src/config.rs", 1, "struct Config { retries: u32 }"),
            ],
            calls: vec![CallEdgeRow {
                crate_name: "fixture".into(),
                caller_path: "retry_fetch".into(),
                callee_ident: "add".into(),
                call_kind: "call".into(),
                file: "src/http.rs".into(),
                line: 12,
            }],
        }
    }

    // ── RAG retrieval ───────────────────────────────────────────────────

    /// inject-assert: a known query retrieves the RELEVANT symbol first. The
    /// "retry the http fetch" query must anchor onto `retry_fetch` in http.rs,
    /// not `add`/`Config` — real input, real asserted ranking.
    #[test]
    fn code_index_retrieval_ranks_relevant_symbol_first() {
        let v = view();
        let r = CodeIndexRetriever::new(&v);
        let hits = r.retrieve("retry the http fetch on failure", 3);
        assert!(!hits.is_empty(), "must retrieve at least one anchor");
        assert_eq!(hits[0].symbol, "retry_fetch", "most relevant symbol first");
        assert_eq!(hits[0].file, "src/http.rs");
        assert_eq!(hits[0].source, "code_index");
        assert!(hits[0].score > 0.0, "real positive relevance score");
    }

    /// A query about math anchors onto `add`, proving the scorer discriminates.
    #[test]
    fn code_index_retrieval_discriminates_by_query() {
        let v = view();
        let r = CodeIndexRetriever::new(&v);
        let hits = r.retrieve("add two integers", 1);
        assert_eq!(hits[0].symbol, "add");
        assert_eq!(hits[0].file, "src/math.rs");
    }

    /// `retrieve` merges + dedups + globally ranks across retrievers.
    #[test]
    fn retrieve_merges_dedups_and_ranks() {
        let v = view();
        let r = CodeIndexRetriever::new(&v);
        // The same retriever twice: anchors must be deduped by (file, symbol).
        let ctx = retrieve(&[&r, &r], "config retries struct", 5);
        let keys: Vec<(String, String)> =
            ctx.anchors.iter().map(|a| (a.file.clone(), a.symbol.clone())).collect();
        let mut uniq = keys.clone();
        uniq.sort();
        uniq.dedup();
        assert_eq!(keys.len(), uniq.len(), "anchors deduped by (file, symbol)");
        assert_eq!(ctx.anchors[0].symbol, "Config", "config query anchors on Config");
        assert_eq!(ctx.retrievers, vec!["code_index"], "only code index contributed");
    }

    // ── decomposition ───────────────────────────────────────────────────

    /// inject-assert: decomposing a task yields N child subtasks each anchored
    /// to a REAL file from the index, capped at `max`, indexed 0..N.
    #[test]
    fn decompose_yields_anchored_subtasks_capped() {
        let v = view();
        let r = CodeIndexRetriever::new(&v);
        let retrievers: Vec<&dyn Retriever> = vec![&r];
        // The query touches all three symbols; cap at 2.
        let d = decompose("rework add and retry_fetch and Config", &retrievers, 5, 2);
        assert_eq!(d.subtasks.len(), 2, "capped at max=2");
        // Every subtask is anchored to a real file in the index.
        let known: std::collections::HashSet<&str> =
            ["src/math.rs", "src/http.rs", "src/config.rs"].into_iter().collect();
        for (i, s) in d.subtasks.iter().enumerate() {
            assert_eq!(s.index, i, "subtasks indexed in order");
            assert!(known.contains(s.target_file.as_str()), "anchored to real file: {}", s.target_file);
            assert!(!s.anchor_symbol.is_empty(), "has a concrete anchor symbol");
            assert!(s.acceptance.contains("cargo"), "has an acceptance check");
            assert_eq!(s.context.anchors.len(), 1, "per-subtask single anchor");
        }
    }

    /// The ≤20 cap holds even when more anchors exist (synthesize 25 symbols).
    #[test]
    fn decompose_respects_max_subtasks_cap() {
        let mut symbols = Vec::new();
        for i in 0..25 {
            symbols.push(sym(
                &format!("widget_{i}"),
                "fn",
                &format!("src/w{i}.rs"),
                1,
                "fn widget()",
            ));
        }
        let v = KnowledgeView { symbols, calls: vec![] };
        let r = CodeIndexRetriever::new(&v);
        let retrievers: Vec<&dyn Retriever> = vec![&r];
        let d = decompose("touch every widget", &retrievers, 100, 100);
        assert!(d.subtasks.len() <= MAX_SUBTASKS, "never exceeds the hard cap");
    }

    /// An empty index still yields one runnable whole-task subtask.
    #[test]
    fn decompose_empty_index_yields_fallback_subtask() {
        let v = KnowledgeView { symbols: vec![], calls: vec![] };
        let r = CodeIndexRetriever::new(&v);
        let retrievers: Vec<&dyn Retriever> = vec![&r];
        let d = decompose("do the thing", &retrievers, 5, 5);
        assert_eq!(d.subtasks.len(), 1, "fallback single subtask");
        assert!(d.subtasks[0].title.contains("do the thing"));
    }

    // ── per-subtask loop with the REAL judge oracle ─────────────────────

    /// A scratch crate whose `src/lib.rs` we overwrite, with baked-in acceptance
    /// tests asserting `add(2,2) == 4`.
    fn scratch_crate() -> tempfile::TempDir {
        let dir = tempfile::tempdir().unwrap();
        let root = dir.path();
        std::fs::write(
            root.join("Cargo.toml"),
            "[package]\nname=\"fix\"\nversion=\"0.0.0\"\nedition=\"2021\"\n[lib]\npath=\"src/lib.rs\"\n",
        )
        .unwrap();
        std::fs::create_dir_all(root.join("src")).unwrap();
        std::fs::write(root.join("src/lib.rs"), "pub fn add(_a:i32,_b:i32)->i32{0}\n").unwrap();
        dir
    }

    const ACCEPT_TESTS: &str = r#"
#[cfg(test)]
mod t { use super::*;
    #[test] fn ok() { assert_eq!(add(2,2), 4); }
}
"#;

    /// A generator that returns a FIXED canned body regardless of the prompt —
    /// lets us drive the loop to a known verdict (good vs non-compiling).
    struct CannedGenerator {
        id: String,
        body: String,
    }
    impl Backend for CannedGenerator {
        fn id(&self) -> &str {
            &self.id
        }
        fn available(&self) -> bool {
            true
        }
    }
    impl Generator for CannedGenerator {
        fn complete(&self, _req: &GenRequest) -> Result<GenAnswer> {
            Ok(GenAnswer {
                text: self.body.clone(),
                tokens_in: 1,
                tokens_out: 1,
                tokens_per_s: 1.0,
                latency_ms: 1.0,
            })
        }
    }

    fn subtask_for(file: &str) -> Subtask {
        Subtask {
            index: 0,
            title: format!("implement add in {file}"),
            target_file: file.into(),
            anchor_symbol: "add".into(),
            acceptance: "cargo build + tests".into(),
            context: RetrievedContext {
                query: "add".into(),
                anchors: vec![],
                retrievers: vec!["code_index".into()],
            },
        }
    }

    /// inject-assert: a KNOWN-GOOD patch is ACCEPTED by the oracle (score 1.0,
    /// node → Done). Real generator → real judge → real cargo build+test.
    #[test]
    fn loop_accepts_known_good_patch() {
        let scratch = scratch_crate();
        let judge = CodegenJudge::new(CodegenSpec::write_file(
            scratch.path().to_path_buf(),
            "src/lib.rs",
        ));
        let good = format!("pub fn add(a:i32,b:i32)->i32{{a+b}}\n{ACCEPT_TESTS}");
        let gen = CannedGenerator { id: "good".into(), body: good };
        let sub = subtask_for("src/lib.rs");
        let v = run_subtask_loop("ship add", &sub, &gen, &judge, 3).unwrap();
        assert!(v.accepted, "good patch accepted: {}", v.best.message);
        assert_eq!(v.best.score, 1.0);
        assert_eq!(v.node_status(), NodeStatus::Done);
        assert_eq!(v.rounds.len(), 1, "accepted on the first round, no retries");
    }

    /// inject-assert: a NON-COMPILING patch is REJECTED (score 0.0, node →
    /// Failed) and the loop exhausts its rounds re-feeding the compiler error.
    #[test]
    fn loop_rejects_noncompiling_patch_and_retries() {
        let scratch = scratch_crate();
        let judge = CodegenJudge::new(CodegenSpec::write_file(
            scratch.path().to_path_buf(),
            "src/lib.rs",
        ));
        // Returns a &str where i32 is expected → never compiles.
        let bad = format!("pub fn add(_a:i32,_b:i32)->i32{{\"nope\"}}\n{ACCEPT_TESTS}");
        let gen = CannedGenerator { id: "bad".into(), body: bad };
        let sub = subtask_for("src/lib.rs");
        let v = run_subtask_loop("ship add", &sub, &gen, &judge, 3).unwrap();
        assert!(!v.accepted, "non-compiling patch rejected");
        assert_eq!(v.best.score, 0.0);
        assert!(!v.best.compiled);
        assert_eq!(v.node_status(), NodeStatus::Failed);
        assert_eq!(v.rounds.len(), 3, "exhausted all bounded rounds on failure");
    }

    /// The retry prompt re-feeds the previous compiler error (the oracle loop).
    #[test]
    fn retry_prompt_carries_prev_error() {
        let sub = subtask_for("src/lib.rs");
        let first = build_codegen_prompt("t", &sub, None);
        assert!(!first.contains("FAILED"), "first round has no error");
        let retry = build_codegen_prompt("t", &sub, Some("error[E0308]: mismatched types"));
        assert!(retry.contains("E0308"), "retry re-feeds the compiler error");
        assert!(retry.contains("Fix it"));
    }

    // ── store subtasks as child funnel DAG nodes ────────────────────────

    /// inject-assert: storing a decomposition creates a child plan whose nodes
    /// mirror the subtasks (anchored targets), chained by edges, only the root
    /// Ready — then applying verdicts flips nodes Done/Failed.
    #[test]
    fn store_subtasks_builds_child_dag_and_applies_verdicts() {
        let dir = tempfile::tempdir().unwrap();
        let mut store = Store::open(dir.path().join("wh")).unwrap();
        let idea = IdeaId::seq(store.funnel.next_idea.max(1));
        store
            .record(Event::IdeaSubmitted {
                id: idea.clone(),
                source: "test".into(),
                text: "rework add and retry_fetch".into(),
                refs: vec![],
                item_kind: ItemKind::Idea,
                ts: Utc::now(),
            })
            .unwrap();

        let v = view();
        let r = CodeIndexRetriever::new(&v);
        let retrievers: Vec<&dyn Retriever> = vec![&r];
        let decomp = decompose("rework add and retry_fetch", &retrievers, 5, 2);
        assert_eq!(decomp.subtasks.len(), 2);

        let (plan_id, node_ids) = store_subtasks(&mut store, &idea, &decomp).unwrap();
        let plan = store.funnel.plans.get(&plan_id).expect("child plan created");
        assert_eq!(plan.idea_id, idea, "child plan hangs under the parent idea");
        assert_eq!(plan.status, PlanStatus::Active);
        assert_eq!(plan.nodes.len(), 2, "one node per subtask");
        assert_eq!(plan.edges.len(), 1, "chained: 2 nodes → 1 edge");

        // Each node targets a real file + carries the subtask state_json (LAW 6).
        for n in plan.nodes.values() {
            assert_eq!(n.kind, "code:write", "idea → code:write child node");
            assert_eq!(n.targets.len(), 1, "node anchored to its file");
            assert!(n.params.contains_key("state_json"), "node emits state_json");
        }

        // Only the root subtask is Ready (chain enforces order).
        store.funnel.promote_ready();
        let plan = store.funnel.plans.get(&plan_id).unwrap();
        let ready: Vec<NodeId> = plan
            .nodes
            .values()
            .filter(|n| n.status == NodeStatus::Ready)
            .map(|n| n.id.clone())
            .collect();
        assert_eq!(ready, vec![node_ids[0].clone()], "only the root subtask is Ready");

        // Apply a mixed verdict set: subtask 0 accepted, subtask 1 failed.
        let verdicts = vec![
            SubtaskVerdict {
                subtask_index: 0,
                accepted: true,
                best: Verdict {
                    compiled: true,
                    tests_passed: 1,
                    tests_total: 1,
                    score: 1.0,
                    message: "all pass".into(),
                },
                rounds: vec![],
            },
            SubtaskVerdict {
                subtask_index: 1,
                accepted: false,
                best: Verdict::rejected("did not compile"),
                rounds: vec![],
            },
        ];
        let n = apply_verdicts(&mut store, &plan_id, &node_ids, &verdicts).unwrap();
        assert_eq!(n, 2, "both verdicts written");
        let plan = store.funnel.plans.get(&plan_id).unwrap();
        assert_eq!(plan.nodes.get(&node_ids[0]).unwrap().status, NodeStatus::Done);
        assert_eq!(plan.nodes.get(&node_ids[1]).unwrap().status, NodeStatus::Failed);
    }

    // ── the offline end-to-end simulation ───────────────────────────────

    /// inject-assert (the headline test): the WHOLE loop runs offline with a
    /// known-good Generator + a REAL CodegenJudge against a fixture crate. The
    /// fixture's one `add` symbol IS the code index → one subtask anchored on
    /// `src/lib.rs`; we drive it with a generator returning a known-good body and
    /// assert the report reflects the oracle verdict end-to-end. No model, no net.
    #[test]
    fn simulate_runs_end_to_end_offline() {
        let scratch = scratch_crate();
        let v = KnowledgeView {
            symbols: vec![sym("add", "fn", "src/lib.rs", 1, "fn add(a:i32,b:i32)->i32")],
            calls: vec![],
        };
        let judge = CodegenJudge::new(CodegenSpec::write_file(
            scratch.path().to_path_buf(),
            "src/lib.rs",
        ));
        let good = format!("pub fn add(a:i32,b:i32)->i32{{a+b}}\n{ACCEPT_TESTS}");
        let gen = CannedGenerator { id: "good".into(), body: good };

        let report = simulate("make add correct", &v, &gen, &judge, 5, 5, 2).unwrap();
        assert_eq!(report.total, 1, "one subtask from the one-symbol index");
        assert_eq!(report.accepted, 1, "the compiler-oracle accepted it");
        assert!(report.all_accepted());
        assert_eq!(report.verdicts[0].best.score, 1.0);
        // The DAG render names the subtask + its PASS verdict.
        let out = report.render();
        assert!(out.contains("PASS"), "render shows the verdict: {out}");
        assert!(out.contains("add"), "render names the anchored symbol");
        // state_json is well-formed and carries the verdict.
        let j = report.to_json();
        assert_eq!(j["accepted"], 1);
        assert_eq!(j["total"], 1);
    }

    /// The `mock` Generator (the offline default) over a crate where the mock's
    /// echo can't compile → the simulation REJECTS (proving the oracle isn't a
    /// rubber stamp). The mock echoes the prompt, which is not valid Rust.
    #[test]
    fn simulate_rejects_when_mock_output_does_not_compile() {
        let scratch = scratch_crate();
        let v = KnowledgeView {
            symbols: vec![sym("add", "fn", "src/lib.rs", 1, "fn add()")],
            calls: vec![],
        };
        let judge = CodegenJudge::new(CodegenSpec::write_file(
            scratch.path().to_path_buf(),
            "src/lib.rs",
        ));
        let gen = MockGenerator::new("mock");
        let report = simulate("make add correct", &v, &gen, &judge, 5, 5, 2).unwrap();
        assert_eq!(report.total, 1);
        assert_eq!(report.accepted, 0, "mock echo is not valid Rust → rejected");
        assert!(!report.all_accepted());
        // It tried the bounded rounds.
        assert_eq!(report.verdicts[0].rounds.len(), 2);
    }
}