//! Real local-model dispatch, metered joules at the meter's honest provenance, in
//! the econ report — and the cheapest-sufficient energy collapse.
//!
//! Wires the pure-Rust `jouleclaw-loader-gguf` inference substrate as a cascade
//! [`Tier`]: it loads a real GGUF/safetensors model, greedily decodes an answer,
//! and meters the energy of that decode window via `jouleclaw-energy`'s
//! auto-detected meter. A recurring task is run several rounds with
//! skill-compilation on: the first few rounds pay the **real model** (real
//! decode time, real metered joules), then the class compiles into a
//! deterministic skill and the rest resolve at ~0 — the cheapest-sufficient
//! thesis in metered energy.
//!
//! HONEST ENERGY: the decode runs on this process's own threads (the loader's CPU
//! kernels on its rayon pool; this example builds the loader without its GPU
//! `fabric`), so the meter is `jouleclaw-energy`'s in-process meter
//! (`detect_in_process_meter`). On macOS that is this process's own CPU-energy
//! counter (`process-cpu`, `Provenance::HwShunt` over `EnergyDomain::ProcessCpu`)
//! where it is live, as it is for an unprivileged reader on macOS 27.0.1 on an M5
//! Max and an M3 Ultra: a hardware reading of the decode's CPU energy, which does
//! not move with anything else the host runs. Where it is not live, and off macOS,
//! it is the host meter (`detect_meter`): the IOReport counter where its CPU rail
//! is live, else a load-average × power-band ESTIMATOR (`Provenance::Estimator`).
//! Each dispatch prints which meter ran and whether its joules were measured,
//! modeled or estimated. Decode wall-time and token counts are exact measurements. A LOCAL
//! model is priced energy-only (no per-call API dollars).
//!
//! Run (defaults to a Nemotron-3-Nano-4B snapshot if present):
//! ```text
//! cargo run --release -p jouleclaw-harness-econ --example local_model_econ
//! cargo run --release -p jouleclaw-harness-econ --example local_model_econ -- /path/to/hf/snapshot
//! ```

use std::path::PathBuf;
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::{joules_obtained, window_joules, EconHarness, PriceModel, Task};
use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::nemotron_h::{NemotronHConfig, NemotronHModel};
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--nvidia--NVIDIA-Nemotron-3-Nano-4B-BF16/snapshots/dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f";

/// One real model dispatch: wall-time, tokens, and metered joules — with the meter
/// that metered them and the provenance of its readings.
#[derive(Clone, Copy)]
struct DispatchStat {
    secs: f64,
    tokens: usize,
    joules: f64,
    meter: &'static str,
    provenance: Provenance,
}

/// A cascade tier backed by a real local model. Greedy-decodes an answer and
/// meters the energy of the decode window. Priced as L3, but LOCAL → the econ
/// PriceModel charges it energy only (no API call dollars).
struct LocalModelTier {
    model: NemotronHModel,
    tok: BpeTokenizer,
    max_tokens: usize,
    bos: u32,
    calls: Arc<AtomicU32>,
    log: Arc<Mutex<Vec<DispatchStat>>>,
}

impl Tier for LocalModelTier {
    fn id(&self) -> TierId {
        TierId::L3(L3ModelId(0))
    }

    fn estimate_cost(&self, _q: &Query) -> Option<TierEstimate> {
        Some(TierEstimate {
            joules: 10.0,
            latency: std::time::Duration::from_secs(1),
            confidence_floor: 0.8,
        })
    }

    fn try_answer(&mut self, q: &Query, _budget: f64) -> Result<Answer, AnswerError> {
        let prompt = match &q.input {
            QueryInput::Text(s) => s.clone(),
            _ => {
                return Err(AnswerError::TierFailed {
                    tier: self.id(),
                    cause: "non-text query".into(),
                })
            }
        };
        let ids: Vec<u32> =
            std::iter::once(self.bos).chain(self.tok.encode(&prompt)).collect();

        // Fresh meter per dispatch (its liveness probe runs once per process): its
        // first read sets the baseline, the second yields the energy over the
        // decode window. 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,
        // a counter that did not advance did not see the work, and a figure the
        // meter did not produce is never reported.
        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)
                }
            });
            generated.push(arg as u32);
        }

        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.calls.fetch_add(1, Ordering::SeqCst);
        self.log.lock().unwrap().push(DispatchStat {
            secs,
            tokens: generated.len(),
            joules,
            meter: meter.active_name(),
            provenance,
        });

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

fn main() {
    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 {}\n(pass an HF snapshot dir with config.json + \
             model.safetensors + tokenizer.json)",
            dir.display()
        );
        return;
    }

    // ---- load the real model ----
    let load = Instant::now();
    let config: NemotronHConfig =
        serde_json::from_slice(&std::fs::read(dir.join("config.json")).unwrap()).unwrap();
    let st = read_safetensors_file(dir.join("model.safetensors"))
        .unwrap_or_else(|e| panic!("safetensors: {e:?}"));
    let model = NemotronHModel::bind(config, st).expect("bind model");
    let tok = BpeTokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer");
    println!("loaded {} in {:.2}s", dir.display(), load.elapsed().as_secs_f64());

    let calls = Arc::new(AtomicU32::new(0));
    let log = Arc::new(Mutex::new(Vec::<DispatchStat>::new()));
    let tier = LocalModelTier {
        model,
        tok,
        max_tokens: 6,
        bos: 1,
        calls: calls.clone(),
        log: log.clone(),
    };
    // run() calls the builder once → take the model out of this slot then.
    let slot = Mutex::new(Some(tier));

    // A recurring factual task the 4B model can answer; verified by substring.
    // (A base 4B model won't write applyable SWE-bench diffs — that's the honest
    // reason this demo uses a task within its reach to exercise the metered tier.)
    let tasks = vec![Task::text(
        "capital-france",
        "geo-capital",
        Query {
            input: QueryInput::Text("The capital of France is".into()),
            budget: JouleBudget::expensive(),
            quality: QualityFloor::any(),
            context: ContextRef::fresh(),
            deadline: None,
        },
        |o| matches!(o, AnswerOutput::Text(t) if t.to_lowercase().contains("paris")),
    )];

    // LOCAL model → energy-only dollars (no per-call API price).
    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 rounds = 5;
    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,
        rounds,
    );

    println!("\n=== Local-model econ report — real inference, metered ===\n");
    print!("{}", report.to_table());

    let stats = log.lock().unwrap();
    println!("\nreal model dispatches: {} (rest served by the compiled skill)", stats.len());
    for (i, s) in stats.iter().enumerate() {
        println!(
            "  dispatch {}: {} tok in {:.2}s ({:.1} tok/s), {:.3} J {} by {} ({})",
            i + 1,
            s.tokens,
            s.secs,
            if s.secs > 0.0 { s.tokens as f64 / s.secs } else { 0.0 },
            s.joules,
            joules_obtained(s.provenance),
            s.meter,
            s.provenance
        );
    }
    let model_j: f64 = stats.iter().map(|s| s.joules).sum();
    let model_n = stats.len().max(1);
    // The dispatches' own provenance, not an assumed one.
    let provenance = match stats.first().map(|s| s.provenance) {
        Some(p) if stats.iter().all(|s| s.provenance == p) => format!("{p:?}"),
        Some(_) => "mixed provenance".to_string(),
        None => "no dispatch".to_string(),
    };
    println!(
        "\navg real-model dispatch  : {:.3} J  ({provenance}, energy-only $)",
        model_j / model_n as f64
    );
    println!(
        "compiled-skill dispatch  : ~{:.0e} J  (deterministic, model untouched)",
        jouleclaw_skill::SKILL_JOULES
    );
    println!(
        "\nverified fixes: {}/{}", report.verified, report.total
    );
    println!(
        "joules per verified fix  : {:.3e} J   ({:?})",
        report.joules_per_verified_fix, report.energy_provenance
    );
    println!("dollars per verified fix : ${:.8}", report.usd_per_verified_fix);
    println!(
        "\nThe energy axis is real inference on this machine, at each dispatch's meter and\n\
         provenance above; the curve bends because the recurring class compiled into a\n\
         skill and stopped paying the model."
    );
}