portal/src/main.rs

//! Generates a portal content repo from a business intake.
//!
//! Three phases, each held to portal's `question_lint`:
//!
//! 1. **Needs.** The intake (what the business does, who comes to it,
//!    what it keeps track of and through which stages, who handles
//!    what) becomes `needs.yaml`. After this phase the file is frozen:
//!    it is the spec, and the generator may not weaken it to pass.
//! 2. **Structure.** Pages, `aggregates.yaml` and desks are generated
//!    and revised against the lint's findings until it exits zero.
//! 3. **Wording.** The personas are answered blind (headings only,
//!    then with descriptions) and misroutes are fed back as wording
//!    defects. Only pages may change.
//!
//! The LLM is an external command: prompt on stdin, text on stdout.
//! `claude -p` works; so does anything else with that shape.

use std::collections::BTreeMap;
use std::io::Write as _;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use anyhow::{anyhow, bail, Context};

mod rate;
mod simulate;

const REFERENCE: &str = include_str!("../generator/reference.md");

struct Opts {
    intake: PathBuf,
    out: PathBuf,
    lint: String,
    llm: String,
    rounds: usize,
    wording_rounds: usize,
    /// The model that plays the visitor. A generator grading its own
    /// wording agrees with itself, so this should be a different one.
    visitor: Option<String>,
}

fn main() -> anyhow::Result<()> {
    // `answer --tasks <file> [--llm <cmd>]`: the blind answerer alone,
    // for grading an existing site with any backend. Answers go to
    // stdout, ready for `question_lint --needs-score`.
    let argv: Vec<String> = std::env::args().skip(1).collect();
    match argv.first().map(String::as_str) {
        Some("simulate") => return simulate::run(&argv),
        Some("compare") => return simulate::compare(&argv),
        Some("rate") => return rate::rate(&argv),
        Some("prefer") => return rate::prefer(&argv),
        _ => {}
    }
    if argv.first().map(String::as_str) == Some("answer") {
        let value = |flag: &str| {
            argv.iter()
                .position(|a| a == flag)
                .and_then(|i| argv.get(i + 1))
                .cloned()
        };
        let tasks = value("--tasks").ok_or_else(|| anyhow!("usage: portal-trainer answer --tasks <tasks.jsonl> [--llm 'cmd']"))?;
        let cmd = value("--llm").unwrap_or_else(|| "claude -p".to_string());
        print!("{}", answer_blind(&cmd, &std::fs::read_to_string(&tasks).with_context(|| format!("reading {tasks}"))?)?);
        return Ok(());
    }
    let opts = parse_args()?;
    let intake = std::fs::read_to_string(&opts.intake)
        .with_context(|| format!("reading {}", opts.intake.display()))?;
    std::fs::create_dir_all(opts.out.join("questions"))?;
    std::fs::create_dir_all(opts.out.join("llm"))?;
    LLM_LOG.set(opts.out.join("llm")).ok();
    let mut log = Vec::new();

    // Phase 1: the spec.
    let mut feedback = String::new();
    for round in 1..=opts.rounds {
        let prompt = format!(
            "{REFERENCE}\n\n# Task\n\nA business filled in this intake. Write ONLY `needs.yaml` for it: every need the intake implies, for the public and for each role inside the business, with states taken from how the business says things progress, and personas.\n\n## Intake\n\n{intake}\n{feedback}\n{FORMAT}"
        );
        let Ok(files) = parse_files(&llm(&opts.llm, &prompt)?) else {
            eprintln!("needs round {round}: the reply had no file blocks, asking again");
            feedback = "\n## Your last attempt\n\nIt contained no <<<FILE ...>>> blocks at all. Return needs.yaml in exactly that format.\n".to_string();
            continue;
        };
        let Some(needs) = files.get("needs.yaml") else {
            feedback = "\n## Your last attempt\n\nIt contained no needs.yaml.\n".to_string();
            continue;
        };
        std::fs::write(opts.out.join("needs.yaml"), needs)?;
        // A placeholder front door, so the lint gets as far as parsing
        // the spec; every need fails, a parse error is what we look for.
        std::fs::write(opts.out.join("questions/index.yaml"), "name: Placeholder?\n")?;
        let (_, report) = lint(&opts, &[])?;
        if let Some(line) = report.lines().find(|l| l.contains("parsing needs.yaml")) {
            eprintln!("needs round {round}: {line}");
            feedback = format!("\n## Your last attempt failed to parse\n\n{line}\n\n```yaml\n{needs}\n```\n");
            continue;
        }
        eprintln!("needs round {round}: needs.yaml parses, frozen");
        log.push(format!("needs: parsed in round {round}"));
        feedback.clear();
        break;
    }
    if !feedback.is_empty() {
        bail!("needs.yaml never parsed in {} round(s); the replies are in {}/llm/", opts.rounds, opts.out.display());
    }
    let needs = std::fs::read_to_string(opts.out.join("needs.yaml")).context("no needs.yaml was produced")?;

    // Phase 2: structure. The page-writing model sees the needs and
    // never the personas: pages written with the test in view route
    // their own personas perfectly and strangers badly (measured:
    // 10/10 on its own personas, 8/10 headings-only on held-out ones).
    let needs_full = needs;
    let needs = needs_full
        .split("\npersonas:")
        .next()
        .unwrap_or(&needs_full)
        .to_string();
    let mut green = false;
    let mut report = String::new();
    for round in 1..=opts.rounds {
        let prompt = if round == 1 {
            format!(
                "{REFERENCE}\n\n# Task\n\nWrite the content repo that meets this frozen `needs.yaml`: `aggregates.yaml`, `site.yaml`, and every page under `questions/` - the front door, followup pages, role-gated pages, one desk per handling group. Do not return needs.yaml.\n\n## Intake\n\n{intake}\n\n## needs.yaml (frozen)\n\n```yaml\n{needs}\n```\n\n{FORMAT}"
            )
        } else {
            format!(
                "{REFERENCE}\n\n# Task\n\nThe verifier rejected the repo below. Fix it. Return only the files you change, each complete. To remove a file, return it with the single line DELETE. needs.yaml is frozen; do not return it.\n\n## Verifier output\n\n```\n{report}\n```\n\n## needs.yaml (frozen)\n\n```yaml\n{needs}\n```\n\n## Current files\n\n{}\n\n{FORMAT}",
                dump(&opts.out)?
            )
        };
        match parse_files(&llm(&opts.llm, &prompt)?) {
            Ok(files) => apply(&opts.out, files)?,
            Err(_) => eprintln!("structure round {round}: the reply had no file blocks"),
        }
        let (ok, r) = lint(&opts, &[])?;
        report = r;
        let fails = report.lines().filter(|l| l.starts_with("FAIL")).count();
        eprintln!("structure round {round}: {}", if ok { "green".to_string() } else { format!("{fails} FAIL line(s)") });
        log.push(format!("structure round {round}: ok={ok} fails={fails}"));
        if ok {
            green = true;
            break;
        }
    }
    if !green {
        std::fs::write(opts.out.join("run.log"), format!("{}\n\n{report}", log.join("\n")))?;
        bail!("structure never went green in {} round(s); last verifier output is in run.log", opts.rounds);
    }

    // Phase 3: wording.
    for round in 0..=opts.wording_rounds {
        let mut defects = String::new();
        let mut summary = Vec::new();
        for skim in [true, false] {
            let tasks_path = opts.out.join(if skim { "tasks-skim.jsonl" } else { "tasks.jsonl" });
            let mut args = vec!["--needs-tasks".to_string(), tasks_path.display().to_string()];
            if skim {
                args.push("--skim".to_string());
            }
            lint(&opts, &args)?;
            let answers = answer_blind(opts.visitor.as_deref().unwrap_or(&opts.llm), &std::fs::read_to_string(&tasks_path)?)?;
            let answers_path = opts.out.join(if skim { "answers-skim.jsonl" } else { "answers.jsonl" });
            std::fs::write(&answers_path, answers)?;
            let mut args = vec!["--needs-score".to_string(), answers_path.display().to_string()];
            if skim {
                args.push("--skim".to_string());
            }
            let (_, scored) = lint(&opts, &args)?;
            let mode = if skim { "headings only" } else { "headings and descriptions" };
            for line in scored.lines().filter(|l| l.starts_with("MISROUTE") || l.starts_with("HESITANT")) {
                defects.push_str(&format!("[{mode}] {line}\n"));
            }
            if let Some(line) = scored.lines().find(|l| l.starts_with("wording:")) {
                summary.push(format!("{mode}: {line}"));
            }
        }
        eprintln!("wording round {round}: {}", summary.join(" | "));
        log.push(format!("wording round {round}: {}", summary.join(" | ")));
        if defects.is_empty() || round == opts.wording_rounds {
            break;
        }
        let personas = needs_full.split("\npersonas:").nth(1).unwrap_or_default();
        let prompt = format!(
            "{REFERENCE}\n\n# Task\n\nThe structure is correct. Visitors are picking the wrong alternative. Each MISROUTE below is a persona who, reading the page, chose an alternative that does not meet their need; HESITANT means they chose right and nearly did not. Rewrite alternative headings first (a heading that needs its description has failed), descriptions second. Keep every `record_as`, `action`, field and state as it is. If you rename an alternative that needs.yaml names in `via`, you may not: keep that one. Return only the pages you change, each complete.\n\n## Defects\n\n```\n{defects}```\n\n## The personas\n\n```yaml\npersonas:{personas}\n```\n\n## Current files\n\n{}\n\n{FORMAT}",
            dump(&opts.out)?
        );
        let before = snapshot(&opts.out)?;
        apply(&opts.out, parse_files(&llm(&opts.llm, &prompt)?)?)?;
        let (ok, r) = lint(&opts, &[])?;
        if !ok {
            eprintln!("wording round {round}: the rewrite broke structure, reverted");
            log.push(format!("wording round {round}: rewrite reverted\n{r}"));
            restore(&opts.out, before)?;
            break;
        }
    }
    std::fs::write(opts.out.join("run.log"), log.join("\n") + "\n")?;
    println!("content repo written to {}", opts.out.display());
    Ok(())
}

const FORMAT: &str = "# Output format\n\nReturn files and nothing else, each as:\n\n<<<FILE relative/path.yaml>>>\n...complete file content...\n<<<END>>>\n\nNo prose, no code fences.";

fn parse_args() -> anyhow::Result<Opts> {
    let mut args = std::env::args().skip(1);
    if args.next().as_deref() != Some("generate") {
        bail!("usage: portal-trainer generate --intake <file> --out <dir> [--lint question_lint] [--llm 'claude -p'] [--rounds 4] [--wording-rounds 2]");
    }
    let mut opts = Opts {
        intake: PathBuf::new(),
        out: PathBuf::new(),
        lint: "question_lint".to_string(),
        llm: "claude -p".to_string(),
        rounds: 4,
        wording_rounds: 2,
        visitor: None,
    };
    while let Some(arg) = args.next() {
        let mut value = || args.next().ok_or_else(|| anyhow!("{arg} needs a value"));
        match arg.as_str() {
            "--intake" => opts.intake = value()?.into(),
            "--out" => opts.out = value()?.into(),
            "--lint" => opts.lint = value()?,
            "--llm" => opts.llm = value()?,
            "--rounds" => opts.rounds = value()?.parse()?,
            "--wording-rounds" => opts.wording_rounds = value()?.parse()?,
            "--visitor-llm" => opts.visitor = Some(value()?),
            other => bail!("unknown argument {other}"),
        }
    }
    if opts.intake.as_os_str().is_empty() || opts.out.as_os_str().is_empty() {
        bail!("--intake and --out are required");
    }
    Ok(opts)
}

/// Prompt on stdin, text on stdout.
/// Where `generate` keeps every prompt and reply (`<out>/llm/`). A
/// model that returns nothing usable has to be read, not guessed at.
static LLM_LOG: std::sync::OnceLock<PathBuf> = std::sync::OnceLock::new();
static LLM_CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);

fn llm(cmd: &str, prompt: &str) -> anyhow::Result<String> {
    let reply = llm_with(cmd, prompt, &[]);
    if let Some(dir) = LLM_LOG.get() {
        let n = LLM_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
        std::fs::write(dir.join(format!("{n:02}-prompt.md")), prompt).ok();
        let text = match &reply {
            Ok(text) => text.clone(),
            Err(e) => format!("ERROR: {e}"),
        };
        std::fs::write(dir.join(format!("{n:02}-reply.md")), text).ok();
    }
    reply
}

/// As `llm`, with hints for the backend in its environment (`SEED`,
/// `FORMAT=json`, `TEMPERATURE`). A backend is free to ignore them.
pub(crate) fn llm_with(cmd: &str, prompt: &str, env: &[(&str, &str)]) -> anyhow::Result<String> {
    let mut child = Command::new("sh")
        .arg("-c")
        .arg(cmd)
        .envs(env.iter().copied())
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
        .with_context(|| format!("starting {cmd:?}"))?;
    let mut stdin = child.stdin.take().expect("piped");
    let prompt = prompt.to_string();
    // Written from a thread: a large prompt would otherwise deadlock
    // against a child that starts writing before it finishes reading.
    let writer = std::thread::spawn(move || stdin.write_all(prompt.as_bytes()));
    let output = child.wait_with_output()?;
    writer.join().ok();
    if !output.status.success() {
        bail!("{cmd:?} exited with {}", output.status);
    }
    Ok(String::from_utf8_lossy(&output.stdout).into_owned())
}

fn lint(opts: &Opts, extra: &[String]) -> anyhow::Result<(bool, String)> {
    let output = Command::new(&opts.lint)
        .arg("--path")
        .arg(opts.out.join("questions"))
        .args(extra)
        .output()
        .with_context(|| format!("running {}", opts.lint))?;
    let mut report = String::from_utf8_lossy(&output.stdout).into_owned();
    report.push_str(&String::from_utf8_lossy(&output.stderr));
    Ok((output.status.success(), report))
}

/// `<<<FILE path>>> ... <<<END>>>` blocks. Anything outside a block is
/// ignored, so a model that adds a sentence does no harm.
fn parse_files(text: &str) -> anyhow::Result<BTreeMap<String, String>> {
    // Small models mangle the markers: `<<<END>>` with a bracket short,
    // a stray `</think>`, the next `<<<FILE` glued onto the same line.
    // So markers are found anywhere in a line, an END needs only its
    // opening, and a new FILE closes whatever was open.
    let mut files = BTreeMap::new();
    let mut current: Option<(String, String)> = None;
    let mut close = |current: &mut Option<(String, String)>, files: &mut BTreeMap<String, String>| {
        if let Some((path, body)) = current.take() {
            files.insert(path, body);
        }
    };
    for line in text.lines() {
        let mut rest = line;
        loop {
            let (file_at, end_at) = (rest.find("<<<FILE "), rest.find("<<<END"));
            match (file_at, end_at) {
                (Some(f), e) if e.is_none_or(|e| f < e) => {
                    close(&mut current, &mut files);
                    let after = &rest[f + "<<<FILE ".len()..];
                    let stop = after.find(">>>").unwrap_or(after.len());
                    current = Some((after[..stop].trim().to_string(), String::new()));
                    rest = after.get(stop + 3..).unwrap_or("");
                }
                (_, Some(e)) => {
                    close(&mut current, &mut files);
                    rest = &rest[e + "<<<END".len()..];
                }
                _ => break,
            }
        }
        // A line that carried a marker carries no content.
        if std::ptr::eq(rest, line) {
            if let Some((_, body)) = current.as_mut() {
                body.push_str(line);
                body.push('\n');
            }
        }
    }
    close(&mut current, &mut files);
    if files.is_empty() {
        bail!("the model returned no <<<FILE ...>>> blocks");
    }
    Ok(files)
}

/// Model output is written to disk, so the path is checked, not
/// trusted: repo-relative, yaml, and only where content lives.
fn allowed(path: &str) -> bool {
    let p = Path::new(path);
    let plain = p.components().all(|c| matches!(c, std::path::Component::Normal(_)));
    plain
        && p.extension().and_then(|e| e.to_str()) == Some("yaml")
        && (matches!(path, "aggregates.yaml" | "site.yaml") || path.starts_with("questions/"))
}

fn apply(out: &Path, files: BTreeMap<String, String>) -> anyhow::Result<()> {
    for (path, body) in files {
        if !allowed(&path) {
            eprintln!("ignored {path:?}: outside what the generator may write");
            continue;
        }
        let target = out.join(&path);
        if body.trim() == "DELETE" {
            std::fs::remove_file(&target).ok();
            continue;
        }
        if let Some(parent) = target.parent() {
            std::fs::create_dir_all(parent)?;
        }
        std::fs::write(target, body)?;
    }
    Ok(())
}

fn content_files(out: &Path) -> anyhow::Result<Vec<PathBuf>> {
    fn walk(dir: &Path, found: &mut Vec<PathBuf>) -> anyhow::Result<()> {
        for entry in std::fs::read_dir(dir)? {
            let path = entry?.path();
            if path.is_dir() {
                walk(&path, found)?;
            } else if path.extension().and_then(|e| e.to_str()) == Some("yaml") {
                found.push(path);
            }
        }
        Ok(())
    }
    let mut found = Vec::new();
    for name in ["aggregates.yaml", "site.yaml"] {
        if out.join(name).exists() {
            found.push(out.join(name));
        }
    }
    walk(&out.join("questions"), &mut found)?;
    found.sort();
    Ok(found)
}

fn snapshot(out: &Path) -> anyhow::Result<Vec<(PathBuf, String)>> {
    content_files(out)?
        .into_iter()
        .map(|p| Ok((p.clone(), std::fs::read_to_string(p)?)))
        .collect()
}

fn restore(out: &Path, before: Vec<(PathBuf, String)>) -> anyhow::Result<()> {
    for path in content_files(out)? {
        std::fs::remove_file(path)?;
    }
    for (path, body) in before {
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        std::fs::write(path, body)?;
    }
    Ok(())
}

fn dump(out: &Path) -> anyhow::Result<String> {
    let mut text = String::new();
    for path in content_files(out)? {
        let rel = path.strip_prefix(out).unwrap_or(&path).display().to_string();
        text.push_str(&format!("<<<FILE {rel}>>>\n{}<<<END>>>\n\n", std::fs::read_to_string(&path)?));
    }
    Ok(text)
}

/// The answering side never sees `expect` or which need a persona has.
fn answer_blind(cmd: &str, tasks: &str) -> anyhow::Result<String> {
    let mut blind = String::new();
    for line in tasks.lines().filter(|l| !l.trim().is_empty()) {
        let mut task: serde_json::Value = serde_json::from_str(line)?;
        if let Some(map) = task.as_object_mut() {
            map.remove("expect");
            map.remove("need");
        }
        blind.push_str(&serde_json::to_string(&task)?);
        blind.push('\n');
    }
    if blind.is_empty() {
        return Ok(String::new());
    }
    let prompt = format!(
        "You are a typed-decision answering engine. Each line below is a JSON task: `state.body` is a website visitor describing their situation in their own words; `questions.pick` holds `instructions` and `criteria`, a map from option label to what that option says.\n\nFor each task decide which single option that visitor would most plausibly choose, judging only from their text and the option wording, and give your honest confidence between 0 and 1. If two options are similarly plausible, confidence is low.\n\nReply with ONLY JSONL, one line per task, same order:\n{{\"id\":\"<task id>\",\"choice\":\"<option label copied exactly>\",\"confidence\":0.00}}\n\n{blind}"
    );
    let reply = llm(cmd, &prompt)?;
    Ok(reply
        .lines()
        .map(str::trim)
        .filter(|l| l.starts_with('{') && serde_json::from_str::<serde_json::Value>(l).is_ok())
        .collect::<Vec<_>>()
        .join("\n")
        + "\n")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn files_are_parsed_from_blocks_and_prose_is_ignored() {
        let files = parse_files("Sure!\n<<<FILE site.yaml>>>\ntitle: A\n<<<END>>>\nnoise\n<<<FILE questions/index.yaml>>>\nname: Q?\n<<<END>>>\n").unwrap();
        assert_eq!(files["site.yaml"], "title: A\n");
        assert_eq!(files.len(), 2);
        assert!(parse_files("no blocks here").is_err());
    }

    #[test]
    fn mangled_markers_from_a_small_model_still_separate_the_files() {
        // Verbatim shape of a nemotron reply: a short END, a stray
        // think tag, the next FILE glued on, and no END at the very end.
        let files = parse_files(
            "<<<FILE questions/desk/index.yaml>>>\nname: Desk?\n<<<END>></think><<<FILE questions/index.yaml>>>\nname: Hva trenger du?\n<<<FILE site.yaml>>>\ntitle: A\n",
        )
        .unwrap();
        assert_eq!(files["questions/desk/index.yaml"], "name: Desk?\n");
        assert_eq!(files["questions/index.yaml"], "name: Hva trenger du?\n");
        assert_eq!(files["site.yaml"], "title: A\n");
    }

    #[test]
    fn only_content_paths_may_be_written() {
        assert!(allowed("site.yaml"));
        assert!(allowed("aggregates.yaml"));
        assert!(allowed("questions/desk/_section.yaml"));
        // The spec is frozen after phase one.
        assert!(!allowed("needs.yaml"));
        assert!(!allowed("../outside.yaml"));
        assert!(!allowed("/etc/passwd.yaml"));
        assert!(!allowed("questions/run.sh"));
        assert!(!allowed(".gitea/workflows/deploy.yaml"));
    }
}