hypersteeldb 0.5.2

//! Shared trajectory-generation logic for the SFT harness (used by the deterministic `gen_trajectories`
//! bin and the Bedrock-teacher `gen_trajectories_teacher` bin). The engine owns correctness — real tool
//! sequences run via `exec_tool`, results are byte-identical to inference — and a teacher (optional) only
//! rewrites the language (natural question + grounded answer) from those real results.

use crate::agent::harness::HARNESS_GUIDANCE;
use crate::agent::tools::{engine_tools, exec_tool, SYSTEM_PROMPT};
use crate::{Corpus, CorpusKind};
use serde_json::{json, Value};

/// One executed tool call in a trajectory (call + the real result JSON string).
pub struct Step {
    pub name: String,
    pub input: Value,
    pub result: String,
}

/// One trajectory: a question, the executed tool chain, and a deterministic gold answer templated from
/// the final program result. A teacher may override `question`/`answer` with natural language.
pub struct Traj {
    pub question: String,
    pub program: String,
    pub steps: Vec<Step>,
    pub result: Value, // parsed last result (for the answer prompt / templating)
    pub answer: String,
}

impl Traj {
    /// A compact one-line digest of the final result — the ground-truth facts a teacher must phrase
    /// without adding anything.
    pub fn facts(&self) -> String {
        self.answer.clone()
    }

    /// Render as an OpenAI `{tools, messages}` SFT example. Optional overrides swap in natural language.
    pub fn to_example(&self, tools: &Value, question: Option<&str>, answer: Option<&str>) -> Value {
        let mut messages: Vec<Value> = vec![
            json!({ "role": "system", "content": format!("{SYSTEM_PROMPT}\n\n{HARNESS_GUIDANCE}") }),
            json!({ "role": "user", "content": question.unwrap_or(&self.question) }),
        ];
        for (i, s) in self.steps.iter().enumerate() {
            let id = format!("call_{i}");
            messages.push(json!({
                "role": "assistant",
                "content": Value::Null,
                "tool_calls": [{ "id": id, "type": "function", "function": { "name": s.name, "arguments": s.input.to_string() } }],
            }));
            messages.push(json!({ "role": "tool", "tool_call_id": id, "content": s.result }));
        }
        messages.push(json!({ "role": "assistant", "content": answer.unwrap_or(&self.answer) }));
        json!({ "tools": tools, "messages": messages })
    }
}

/// The OpenAI-format tool definitions (same for every example).
pub fn tools_json() -> Value {
    json!(engine_tools(&[])
        .iter()
        .map(|t| json!({ "type": "function", "function": { "name": t.name, "description": t.description, "parameters": t.schema } }))
        .collect::<Vec<_>>())
}

// ── synthetic corpora with randomized schemas (variety for generalization) ──────────

struct Domain {
    name: &'static str,
    facets: &'static [(&'static str, &'static [&'static str])],
}

const DOMAINS: &[Domain] = &[
    Domain { name: "vehicles", facets: &[("make", &["toyota", "honda", "hyundai", "kia", "bmw", "tesla", "ford", "mazda", "volvo"]), ("country", &["japan", "korea", "germany", "usa", "sweden"]), ("body", &["sedan", "suv", "hatchback", "truck", "coupe"]), ("powertrain", &["electric", "hybrid", "petrol", "diesel"])] },
    Domain { name: "incidents", facets: &[("service", &["auth", "payments", "search", "checkout", "notifications", "billing"]), ("severity", &["low", "medium", "high", "critical"]), ("region", &["us", "eu", "apac", "latam"]), ("status", &["open", "mitigated", "resolved"]), ("cause", &["deploy", "capacity", "dependency", "config"])] },
    Domain { name: "papers", facets: &[("field", &["ml", "systems", "theory", "security", "graphics", "hci"]), ("venue", &["neurips", "osdi", "sigcomm", "focs", "siggraph", "chi"]), ("outcome", &["accepted", "rejected", "spotlight", "poster"]), ("method", &["empirical", "theoretical", "survey"])] },
    Domain { name: "products", facets: &[("category", &["laptop", "phone", "tablet", "watch", "headphones"]), ("brand", &["apple", "samsung", "dell", "sony", "lenovo"]), ("tier", &["budget", "mid", "premium", "flagship"]), ("availability", &["instock", "backorder", "discontinued"])] },
    Domain { name: "properties", facets: &[("type", &["apartment", "house", "condo", "townhouse"]), ("city", &["seattle", "austin", "denver", "miami", "boston"]), ("status", &["listed", "pending", "sold"]), ("heating", &["gas", "electric", "heatpump"])] },
    Domain { name: "recipes", facets: &[("cuisine", &["italian", "japanese", "mexican", "indian", "thai"]), ("course", &["appetizer", "main", "dessert", "soup"]), ("diet", &["vegan", "vegetarian", "omnivore", "glutenfree"]), ("difficulty", &["easy", "medium", "hard"])] },
    Domain { name: "flights", facets: &[("carrier", &["delta", "united", "ana", "lufthansa", "emirates"]), ("origin", &["sea", "sfo", "jfk", "nrt", "lhr"]), ("cabin", &["economy", "premium", "business", "first"]), ("status", &["ontime", "delayed", "cancelled"])] },
    Domain { name: "grants", facets: &[("agency", &["nsf", "nih", "darpa", "doe"]), ("area", &["biology", "physics", "cs", "materials"]), ("stage", &["submitted", "review", "awarded", "declined"]), ("size", &["small", "medium", "large"])] },
];

/// Build `n` synthetic corpora with randomized schemas (a random domain, a random subset of its facets,
/// a random subset of values, and a random row count). Seeded for reproducibility.
pub fn synth_corpora(n: usize, seed: u32) -> Vec<(String, Corpus)> {
    let mut s = seed;
    let mut rnd = move |m: usize| {
        s = s.wrapping_mul(1664525).wrapping_add(1013904223);
        (s >> 8) as usize % m.max(1)
    };
    let mut out = Vec::new();
    for i in 0..n {
        let dom = &DOMAINS[rnd(DOMAINS.len())];
        // pick 2-4 facets
        let nf = 2 + rnd(3.min(dom.facets.len() - 1));
        let mut fidx: Vec<usize> = (0..dom.facets.len()).collect();
        for j in 0..fidx.len() {
            let k = rnd(fidx.len());
            fidx.swap(j, k);
        }
        fidx.truncate(nf);
        let chosen: Vec<(&str, Vec<&str>)> = fidx
            .iter()
            .map(|&fi| {
                let (fname, vals) = dom.facets[fi];
                // pick a random subset (>=2) of values
                let take = 2 + rnd(vals.len() - 1);
                let mut idx: Vec<usize> = (0..vals.len()).collect();
                for j in 0..idx.len() {
                    let k = rnd(idx.len());
                    idx.swap(j, k);
                }
                idx.truncate(take.min(vals.len()));
                (fname, idx.into_iter().map(|k| vals[k]).collect())
            })
            .collect();

        let cols: Vec<String> = chosen.iter().map(|(f, _)| f.to_string()).collect();
        let cname = format!("{}-{i}", dom.name);
        let mut c = Corpus::new_incremental(cname.clone(), cols, CorpusKind::Csv);
        let rows = 40 + rnd(80);
        for _ in 0..rows {
            let mut tokens = Vec::new();
            let mut cells = Vec::new();
            for (f, vals) in &chosen {
                let v = vals[rnd(vals.len())];
                tokens.push(format!("{f}/{v}"));
                cells.push(v.to_string());
            }
            c.add_situation(tokens, cells);
        }
        out.push((cname, c));
    }
    out
}

// ── deterministic trajectory templates (exercise every bitmap program) ───────────────

fn run(corpus: &Corpus, plan: &[(&str, Value)]) -> (Vec<Step>, Value) {
    let mut steps = Vec::new();
    let mut last = Value::Null;
    for (name, input) in plan {
        let (result, _e) = exec_tool(corpus, name, input);
        last = serde_json::from_str(&result).unwrap_or(Value::Null);
        steps.push(Step { name: name.to_string(), input: input.clone(), result });
    }
    (steps, last)
}

pub fn gen_templated(corpus: &Corpus) -> Vec<Traj> {
    let facets: Vec<(String, usize)> = corpus.stats().facets.into_iter().filter(|(f, n)| *n >= 2 && *n <= 20 && f != "src").collect();
    let mut out = Vec::new();
    if facets.is_empty() {
        return out;
    }

    for (f, _) in &facets {
        let (steps, r) = run(corpus, &[("list_facets", json!({})), ("facet_tokens", json!({ "facet": f })), ("breakdown", json!({ "anchor": "*", "facet": f }))]);
        out.push(Traj { question: format!("How does the corpus break down by {f}?"), program: "breakdown".into(), answer: ans_breakdown(&r), steps, result: r });

        let (steps, r) = run(corpus, &[("list_facets", json!({})), ("rank", json!({ "facet": f }))]);
        out.push(Traj { question: format!("Which {f} values are most significant?"), program: "rank".into(), answer: ans_rank(&r), steps, result: r });
    }

    for pair in facets.windows(2).take(3) {
        let (a, b) = (pair[0].0.clone(), pair[1].0.clone());
        let (steps, r) = run(corpus, &[("list_facets", json!({})), ("facet_tokens", json!({ "facet": &a })), ("facet_tokens", json!({ "facet": &b })), ("crosstab", json!({ "anchor": "*", "facet_a": &a, "facet_b": &b }))]);
        out.push(Traj { question: format!("How do {a} and {b} relate across the corpus?"), program: "crosstab".into(), answer: ans_crosstab(&r), steps, result: r });
    }

    // FILTER-THEN-AGGREGATE — the multi-constraint shape ("which G has the most where F=v"): put the
    // filter token in the anchor and break down by the OTHER facet. This is what a whole-corpus-only
    // training set never teaches, so it's essential coverage for real analytic questions.
    for pair in facets.windows(2).take(4) {
        let (f, g) = (pair[0].0.clone(), pair[1].0.clone());
        if let Some((tok, _)) = corpus.facet_tokens(&f, 1).into_iter().next() {
            let v = tok.split('/').nth(1).unwrap_or(&tok).to_string();
            let (steps, r) = run(corpus, &[("list_facets", json!({})), ("facet_tokens", json!({ "facet": &f })), ("facet_tokens", json!({ "facet": &g })), ("breakdown", json!({ "anchor": &tok, "facet": &g }))]);
            let ans = ans_filtered(&r, &f, &v, &g);
            out.push(Traj { question: format!("For {f} = {v}, how do {g} break down? / Which {g} is most common among {f} {v}?"), program: "breakdown".into(), answer: ans, steps, result: r });
        }
    }

    if let Some((f, _)) = facets.first() {
        if let Some((tok, _)) = corpus.facet_tokens(f, 1).first() {
            let leaf = tok.split('/').nth(1).unwrap_or(tok).to_string();
            let (steps, r) = run(corpus, &[("facet_tokens", json!({ "facet": f })), ("ikl_query", json!({ "ikl": tok }))]);
            out.push(Traj { question: format!("How many situations have {f} = {leaf}?"), program: "ikl_query".into(), answer: ans_ikl(&r), steps, result: r });
        }
    }

    if facets.len() >= 2 {
        let f = &facets[1].0;
        if let Some((tok, _)) = corpus.facet_tokens(f, 1).first() {
            let (steps, r) = run(corpus, &[("facet_tokens", json!({ "facet": f })), ("cooccurs", json!({ "token": tok }))]);
            out.push(Traj { question: format!("What concepts most co-occur with {tok}?"), program: "cooccurs".into(), answer: ans_cooccurs(&r), steps, result: r });
        }
    }

    if let Some((f, _)) = facets.first() {
        if let Some((real, _)) = corpus.facet_tokens(f, 1).first() {
            let (steps, r) = run(corpus, &[("narrow", json!({ "scope": [real], "filters": ["nonexistent/xyz"] }))]);
            out.push(Traj { question: format!("Are there any {f} situations also tagged nonexistent/xyz?"), program: "narrow".into(), answer: ans_narrow(&r), steps, result: r });
        }
    }

    out
}

fn leaf(t: &str) -> &str {
    t.split('/').nth(1).unwrap_or(t)
}

pub fn ans_breakdown(r: &Value) -> String {
    let facet = r["facet"].as_str().unwrap_or("");
    let total = r["total"].as_u64().unwrap_or(0);
    let parts: Vec<String> = r["partition"].as_array().map(|a| a.iter().take(8).map(|p| format!("{} ({})", p["value"].as_str().unwrap_or(""), p["count"].as_u64().unwrap_or(0))).collect()).unwrap_or_default();
    if parts.is_empty() {
        return format!("No {facet} values are present in the corpus.");
    }
    format!("Across {total} situations, the {facet} breakdown is: {}.", parts.join(", "))
}
/// Answer for a filtered breakdown: "Among F=v, the G breakdown is …; most common is X (n)."
pub fn ans_filtered(r: &Value, f: &str, v: &str, g: &str) -> String {
    let total = r["total"].as_u64().unwrap_or(0);
    let parts: Vec<(String, u64)> = r["partition"]
        .as_array()
        .map(|a| a.iter().take(8).map(|p| (p["value"].as_str().unwrap_or("").to_string(), p["count"].as_u64().unwrap_or(0))).collect())
        .unwrap_or_default();
    if parts.is_empty() {
        return format!("No {g} values occur among {f} {v}.");
    }
    let top = &parts[0];
    let list = parts.iter().map(|(val, n)| format!("{val} ({n})")).collect::<Vec<_>>().join(", ");
    format!("Among the {total} {f}/{v} situations, {g} breaks down as: {list}. The most common is {} ({}).", top.0, top.1)
}

pub fn ans_rank(r: &Value) -> String {
    let facet = r["facet"].as_str().unwrap_or("");
    let toks: Vec<String> = r["ranked"].as_array().map(|a| a.iter().take(5).map(|x| leaf(x["token"].as_str().unwrap_or("")).to_string()).collect()).unwrap_or_default();
    if toks.is_empty() {
        return format!("The corpus has no ranked {facet} values.");
    }
    format!("By salience, the most significant {facet} values are: {}.", toks.join(", "))
}
pub fn ans_crosstab(r: &Value) -> String {
    let (a, b) = (r["row_facet"].as_str().unwrap_or(""), r["col_facet"].as_str().unwrap_or(""));
    let empty = Vec::new();
    let mut lines = Vec::new();
    for row in r["matrix"].as_array().unwrap_or(&empty).iter().take(6) {
        let label = row["row"].as_str().unwrap_or("");
        let cells = row["cells"].as_array().cloned().unwrap_or_default();
        if let Some(t) = cells.iter().max_by_key(|c| c["count"].as_u64().unwrap_or(0)) {
            lines.push(format!("{label} → mostly {} ({})", t["col"].as_str().unwrap_or(""), t["count"].as_u64().unwrap_or(0)));
        }
    }
    if lines.is_empty() {
        return format!("There is no co-occurrence between {a} and {b}.");
    }
    format!("Cross-tabulating {a} × {b}: {}.", lines.join("; "))
}
pub fn ans_ikl(r: &Value) -> String {
    format!("{} situations match `{}`.", r["count"].as_u64().unwrap_or(0), r["ikl"].as_str().unwrap_or(""))
}
pub fn ans_cooccurs(r: &Value) -> String {
    let focus = r["focus"].as_str().unwrap_or("");
    let cs: Vec<String> = r["cooccurs"].as_array().map(|a| a.iter().take(6).map(|x| format!("{} ({})", leaf(x["token"].as_str().unwrap_or("")), x["shared"].as_u64().unwrap_or(0))).collect()).unwrap_or_default();
    if cs.is_empty() {
        return format!("Nothing co-occurs with {focus}.");
    }
    format!("The concepts most co-occurring with {focus} are: {}.", cs.join(", "))
}
pub fn ans_narrow(r: &Value) -> String {
    if r["verdict"].as_str() == Some("ANSWERABLE") {
        let n = r["steps"].as_array().and_then(|s| s.last()).and_then(|s| s["remaining"].as_u64()).unwrap_or(0);
        format!("Yes — {n} situations satisfy all of those constraints together.")
    } else {
        format!("No. Adding `{}` drops the matching set to zero, so the corpus has no supporting evidence.", r["empty_at"].as_str().unwrap_or("that constraint"))
    }
}