//! A varied multi-instance joules/$-per-verified-fix benchmark on a real local
//! code model. Five distinct bug classes (wrong-operator, off-by-one,
//! empty-guard, boundary-comparison, identity-vs-abs) — each a real Python file
//! the Granite coder must fix as a whole-file replacement, judged by a real
//! `python3` test ([`WholeFileValidator`]), with the decode energy metered.
//!
//! This is the multi-instance generalisation of `swebench_local_coder`: instead
//! of one fix it reports a metered curve across bug types — which classes a 4B
//! local model resolves, and the joules/$ each verified fix actually cost. The
//! denominator is fixes that PASS their test, never attempts. Energy covers the decode alone,
//! which runs on this process's own threads (the loader's CPU kernels; built
//! without its GPU `fabric`), so the meter is `jouleclaw-energy`'s in-process meter
//! (`detect_in_process_meter`) at its own provenance: this process's own
//! CPU-energy counter (`process-cpu`, `HwShunt` over `ProcessCpu`) where it is
//! live, as on macOS 27.0.1 on an M5 Max and an M3 Ultra; else the host meter —
//! the IOReport counter (`HwShunt`) where its CPU rail is live, or the
//! load-derived estimator (`Estimator`). The `python3` judge is a subprocess and
//! runs outside the metered window. The report says which meter ran (decode time
//! + tokens are exact). A local model is priced energy-only.
//!
//! Run (defaults to the granite-4.0-h-micro snapshot if present):
//! ```text
//! cargo run --release -p jouleclaw-harness-econ --example swebench_local_suite
//! cargo run --release -p jouleclaw-harness-econ --example swebench_local_suite -- /path/to/snapshot
//! ```

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use jouleclaw_cascade::{
    Answer, AnswerError, AnswerOutput, Cascade, ContextRef, ExecutionTrace, JouleBudget,
    L3ModelId, QualityFloor, Query, QueryInput, Runtime, Tier, TierEstimate, TierId,
    VerificationStatus,
};
use jouleclaw_energy::meter::detect_in_process_meter;
use jouleclaw_energy::Provenance;
use jouleclaw_harness_econ::swebench::exec::{ShellTestRunner, WholeFileValidator};
use jouleclaw_harness_econ::swebench::{PatchValidator, SweInstance};
use jouleclaw_harness_econ::{joules_obtained, window_joules, EconHarness, PriceModel, Task};
use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::granite_h::{GraniteHConfig, GraniteHModel};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use jouleclaw_skill::{ConstantSkillInducer, SkillTier};

const DEFAULT_SNAPSHOT: &str = "/Volumes/macos_4TB_external_1/hf_hub/models--ibm-granite--granite-4.0-h-micro/snapshots/d5f01a3ea75f088947be3aae039f4ad52837dfde";

/// One bug instance: a module name, its buggy source, a one-line description of
/// the fix, and the python assertion that judges it.
struct Bug {
    id: &'static str,
    bug_desc: &'static str,
    buggy: &'static str,
    assertion: &'static str,
}

const BUGS: &[Bug] = &[
    Bug {
        id: "calc",
        bug_desc: "`add` must return the sum a + b, but it returns a - b",
        buggy: "def add(a, b):\n    return a - b\n",
        assertion: "import calc; assert calc.add(2,3)==5 and calc.add(10,4)==14",
    },
    Bug {
        id: "rng",
        bug_desc: "`upto(n)` must return [0..n-1] (n elements) but drops the last via range(n-1)",
        buggy: "def upto(n):\n    return list(range(n - 1))\n",
        assertion: "import rng; assert rng.upto(3)==[0,1,2] and rng.upto(1)==[0]",
    },
    Bug {
        id: "first",
        bug_desc: "`first(xs)` must return None on empty input, currently raises IndexError",
        buggy: "def first(xs):\n    return xs[0]\n",
        assertion: "import first; assert first.first([5,6])==5 and first.first([]) is None",
    },
    Bug {
        id: "adult",
        bug_desc: "`is_adult(age)` must be True at exactly 18 (>=), currently uses >",
        buggy: "def is_adult(age):\n    return age > 18\n",
        assertion: "import adult; assert adult.is_adult(18) and not adult.is_adult(17)",
    },
    Bug {
        id: "absval",
        bug_desc: "`my_abs(x)` must return the absolute value, currently returns x unchanged",
        buggy: "def my_abs(x):\n    return x\n",
        assertion: "import absval; assert absval.my_abs(-5)==5 and absval.my_abs(3)==3",
    },
];

struct DispatchStat {
    secs: f64,
    tokens: usize,
    joules: f64,
    /// The meter that metered the decode, and the provenance of its readings.
    meter: &'static str,
    provenance: Provenance,
    text: String,
}

/// Granite-backed coder tier: chat-template the problem, greedy-decode the
/// corrected file until end-of-text, meter the decode window.
struct LocalCoderTier {
    model: GraniteHModel,
    tok: BpeTokenizer,
    max_tokens: usize,
    eot: Option<u32>,
    log: Arc<Mutex<Vec<DispatchStat>>>,
}

impl Tier for LocalCoderTier {
    fn id(&self) -> TierId {
        TierId::L3(L3ModelId(0))
    }
    fn estimate_cost(&self, _q: &Query) -> Option<TierEstimate> {
        Some(TierEstimate {
            joules: 30.0,
            latency: std::time::Duration::from_secs(20),
            confidence_floor: 0.8,
        })
    }
    fn try_answer(&mut self, q: &Query, _budget: f64) -> Result<Answer, AnswerError> {
        let user = match &q.input {
            QueryInput::Text(s) => s.clone(),
            _ => return Err(AnswerError::TierFailed { tier: self.id(), cause: "non-text".into() }),
        };
        let prompt = format!(
            "<|start_of_role|>user<|end_of_role|>{user}<|end_of_text|>\n\
             <|start_of_role|>assistant<|end_of_role|>"
        );
        let ids = self.tok.encode(&prompt);

        // The decode runs on this process's threads, so the in-process meter
        // covers it. A window the meter put no figure on stops the run
        // (`window_joules`): a failed read is not a zero-energy reading, and a
        // counter that did not advance did not see the work.
        let meter = detect_in_process_meter();
        let e0 = meter.read();
        let t = Instant::now();

        let mut state = self.model.new_state();
        let mut logits = self.model.forward_prompt(&ids, &mut state);
        let mut generated: Vec<u32> = Vec::new();
        for step in 0..self.max_tokens {
            let l = if step == 0 {
                std::mem::take(&mut logits)
            } else {
                self.model.forward_token(*generated.last().unwrap(), &mut state)
            };
            let (arg, _) = l.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
                if v > bv { (i, v) } else { (bi, bv) }
            });
            let tok = arg as u32;
            if Some(tok) == self.eot {
                break;
            }
            generated.push(tok);
        }

        let secs = t.elapsed().as_secs_f64();
        let (joules, provenance) = window_joules(e0, meter.read())
            .unwrap_or_else(|why| panic!("no energy figure for this dispatch from `{}`: {why}", meter.active_name()));
        let text = self.tok.decode(&generated);
        self.log.lock().unwrap().push(DispatchStat {
            secs,
            tokens: generated.len(),
            joules,
            meter: meter.active_name(),
            provenance,
            text: text.clone(),
        });

        Ok(Answer {
            output: AnswerOutput::Text(text),
            tier_used: self.id(),
            joules_spent: joules,
            confidence: 0.8,
            trace: ExecutionTrace::default(),
            verification: VerificationStatus::Resolved,
        })
    }
}

fn git(dir: &Path, args: &[&str]) -> std::process::Output {
    Command::new("git").current_dir(dir).args(args).output().expect("git runs")
}

fn main() {
    if Command::new("python3").arg("--version").output().is_err() {
        eprintln!("python3 not found — this benchmark runs real python tests");
        return;
    }
    let dir = PathBuf::from(std::env::args().nth(1).unwrap_or_else(|| DEFAULT_SNAPSHOT.into()));
    if !dir.join("config.json").is_file() {
        eprintln!("model snapshot not found at {}", dir.display());
        return;
    }

    // ---- build one repo with all buggy files at the base commit ----
    static C: AtomicU32 = AtomicU32::new(0);
    let repo = std::env::temp_dir()
        .join(format!("jc-suite-{}-{}", std::process::id(), C.fetch_add(1, Ordering::SeqCst)));
    std::fs::create_dir_all(&repo).unwrap();
    git(&repo, &["init", "-q"]);
    git(&repo, &["config", "user.email", "t@t.test"]);
    git(&repo, &["config", "user.name", "t"]);
    for b in BUGS {
        std::fs::write(repo.join(format!("{}.py", b.id)), b.buggy).unwrap();
    }
    git(&repo, &["add", "-A"]);
    git(&repo, &["commit", "-q", "-m", "base (buggy)"]);
    let base = String::from_utf8(git(&repo, &["rev-parse", "HEAD"]).stdout).unwrap().trim().to_string();

    // one test runner over all instances' assertions, keyed by node id
    let cmds: BTreeMap<String, String> = BUGS
        .iter()
        .map(|b| (format!("test::{}", b.id), format!("python3 -c '{}'", b.assertion)))
        .collect();
    let runner = ShellTestRunner::new(move |node: &str| cmds.get(node).cloned().unwrap_or_else(|| "false".into()));
    // target file is "<instance_id>.py"
    let validator: Arc<dyn PatchValidator> = Arc::new(WholeFileValidator::new(
        &repo,
        |inst: &SweInstance| PathBuf::from(format!("{}.py", inst.instance_id)),
        runner,
    ));

    // ---- instances + tasks ----
    let mut tasks = Vec::new();
    for b in BUGS {
        let problem = format!(
            "Fix the bug in the Python file `{0}.py`: {1}. Output ONLY the complete corrected \
             contents of {0}.py inside a single ```python code block.\n\n{0}.py:\n```python\n{2}```",
            b.id, b.bug_desc, b.buggy
        );
        let inst = SweInstance {
            instance_id: b.id.into(),
            repo: "local/suite".into(),
            base_commit: base.clone(),
            problem_statement: problem.clone(),
            patch: String::new(),
            test_patch: String::new(),
            fail_to_pass: vec![format!("test::{}", b.id)],
            pass_to_pass: vec![],
            version: None,
            environment_setup_commit: None,
        };
        let v = validator.clone();
        let inst_v = inst.clone();
        tasks.push(Task::text(
            b.id,
            b.id, // each its own class (one-shot, no skill compilation)
            Query {
                input: QueryInput::Text(problem),
                // Budget must cover a real ~200-token local decode (~hundreds of
                // J) so no dispatch is truncated/rejected mid-flight — every
                // attempt's energy is then counted, including failed ones.
                budget: JouleBudget { hard_limit: 2000.0, soft_target: 200.0 },
                quality: QualityFloor::any(),
                context: ContextRef::fresh(),
                deadline: None,
            },
            move |o| matches!(o, AnswerOutput::Text(t) if v.validate(&inst_v, t)),
        ));
    }

    // ---- load the real code model ----
    let load = Instant::now();
    let config: GraniteHConfig =
        serde_json::from_slice(&std::fs::read(dir.join("config.json")).unwrap()).unwrap();
    let mut shard_paths: Vec<PathBuf> = std::fs::read_dir(&dir)
        .unwrap()
        .filter_map(|e| e.ok().map(|e| e.path()))
        .filter(|p| {
            p.extension().is_some_and(|e| e == "safetensors")
                && p.file_name().is_some_and(|n| n.to_string_lossy().starts_with("model"))
        })
        .collect();
    shard_paths.sort();
    let shards: Vec<_> = shard_paths
        .iter()
        .map(|p| read_safetensors_file(p).unwrap_or_else(|e| panic!("{p:?}: {e:?}")))
        .collect();
    let model = GraniteHModel::bind_sharded(config, shards).expect("bind");
    let tok = BpeTokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer");
    let eot = tok.special_id("<|end_of_text|>");
    println!("loaded granite in {:.2}s", load.elapsed().as_secs_f64());

    let log = Arc::new(Mutex::new(Vec::<DispatchStat>::new()));
    let tier = LocalCoderTier { model, tok, max_tokens: 220, eot, log: log.clone() };
    let slot = Mutex::new(Some(tier));

    let mut harness = EconHarness::new(Box::new(ConstantSkillInducer), 3);
    harness.price = PriceModel { usd_per_kwh: 0.15, l3_call_usd: 0.0, l4_call_usd: 0.0 };

    let report = harness.run(
        move |store| {
            let mut casc = Cascade::new();
            casc.register(Box::new(SkillTier::new(store)));
            casc.register(Box::new(slot.lock().unwrap().take().expect("built once")));
            Runtime::new_without_l0(casc)
        },
        &tasks,
        1,
    );

    // ---- report ----
    println!("\n=== Multi-instance local-model econ benchmark — real Granite, real python tests ===\n");
    let stats = log.lock().unwrap();
    println!("{:<8} {:>8} {:>7} {:>9}  {}", "bug", "tokens", "secs", "joules", "result");
    for (i, b) in BUGS.iter().enumerate() {
        let s = &stats[i];
        let resolved = report.rows.iter().find(|r| r.task_id == b.id).map(|r| r.verified).unwrap_or(false);
        println!(
            "{:<8} {:>8} {:>7.1} {:>9.3}  {}",
            b.id,
            s.tokens,
            s.secs,
            s.joules,
            if resolved { "RESOLVED" } else { "unresolved" }
        );
    }
    println!(
        "\nresolved (python tests really ran): {}/{}",
        report.verified, report.total
    );
    let model_j: f64 = stats.iter().map(|s| s.joules).sum();
    let toks: usize = stats.iter().map(|s| s.tokens).sum();
    let secs: f64 = stats.iter().map(|s| s.secs).sum();
    println!(
        "totals: {toks} tok in {secs:.1}s ({:.1} tok/s), {model_j:.3} J spent",
        if secs > 0.0 { toks as f64 / secs } else { 0.0 }
    );
    for s in stats.iter() {
        println!("  {:.3} J {} by {} ({})", s.joules, joules_obtained(s.provenance), s.meter, s.provenance);
    }
    if report.verified > 0 {
        println!(
            "joules per VERIFIED fix : {:.3} J   ({:?})",
            report.joules_per_verified_fix, report.energy_provenance
        );
        println!("dollars per VERIFIED fix: ${:.8}  (local → energy-only)", report.usd_per_verified_fix);
    } else {
        println!("0 verified — honest cascade signal: this tier is insufficient here, escalate.");
    }
    println!(
        "\nThe denominator is fixes that PASS their test, not attempts. Energy is real\n\
         inference on this machine, at each dispatch's meter and provenance above; swap a\n\
         stronger model tier to raise the resolve rate and lower the joules-per-verified-fix."
    );

    drop(stats);
    let _ = std::fs::remove_dir_all(&repo);
}