//! `wai_prov` — author a content-addressed asset + its provenance
//! receipt (lever 4). Produces a real opaque asset (a genuine JPEG, which
//! WAI does NOT reconstruct — it carries it) wrapped in a WAI envelope,
//! and a signed, energy-accounted lineage: origin → PNG master → published
//! JPEG, with the bracketed CPU energy of each encode (the processor power
//! the non-privileged reader `macmon` reports, no sudo), each figure labelled
//! `ModelBased`: the power is one the reader reports from a channel it
//! selects, not a counter this tool reads, so it is never `OnChipCounter`
//! (energy-measurement §3.1; `wai::energy_burst::PowerSource`). An encode whose
//! bracket the reported power's liveness (§3.1) or the §4 ratio cannot support
//! is recorded unmetered: the step's `joules_micro` carries the unmetered marker `0`, with
//! no label, never a figure of zero joules.
//!
//!   wai_prov <out_dir> [--no-report] [--previous-report <file>] [--carry <receipt.json>]
//!
//! Writes `<out_dir>/captured.jpeg.wai` (the content-addressed asset),
//! `<out_dir>/captured.receipt.json` (the signed provenance), and for each
//! measured step `<out_dir>/captured.step<i>.report.json`: its measurement
//! report (energy-measurement §6), a `measurement` claim signed with the
//! receipt's key naming the step by its index, the first chained to
//! `--previous-report` when it is given and the second to the first.
//! `--no-report` turns the reports off, and says so.
//!
//! `--carry <receipt.json>` measures nothing: each step's figure, and its label
//! if it has one, is carried as an earlier receipt for this same lineage states
//! it, and signed under this revision's signer id. That receipt must verify, and
//! its steps must be this run's — the same operations, input and output hashes
//! and work — so a figure is only ever carried onto the bytes it was measured
//! on. No report is written: a report restates a bracket, and this run took
//! none. Nothing in such a run varies, so a committed asset and receipt
//! (`wai-web/reach/test/vectors/provenance/`) are written again byte for byte
//! when the receipt is carried onto itself; a host whose synthesis or encoders
//! give other bytes has another lineage, and is refused. The asset is
//! synthesized deterministically here — WAI is asset-agnostic; the lineage
//! mechanism is identical whether the origin is a camera or a synthesizer,
//! and that is stated plainly. The point is provenance for bytes WAI
//! cannot rebuild, not the bytes' origin.

use ed25519_dalek::SigningKey;
use std::path::Path;
use wai::codecs::image::{jpeg_encode, png_encode};
use wai::energy_burst::{measure_with_macmon, power_reader};
use wai::energy_class::{CounterFiltering, EnergyProvenance};
use wai::energy_figure::meter_report;
use wai::energy_report::BracketValues;
use wai::provenance::{AssetReceipt, ProvStep};
use wai::world::asset_content_hash;
use wai::{Conditioning, Manifest, ModelRequirement, Wai};

const W: u32 = 512;
const H: u32 = 512;

fn main() {
    let mut argv: Vec<String> = std::env::args().collect();
    let carry = take_carry(&mut argv);
    let report_args = meter_report::parse_args_or_exit(&argv);
    let args = &report_args.positional;
    let dir = Path::new(args.get(1).map(String::as_str).unwrap_or("."));
    let carried = carry.as_deref().map(read_carried);
    std::fs::create_dir_all(dir).expect("mkdir");

    // A deterministic synthesized frame (radial rings + gradient). Real
    // pixels, no clock — the lineage mechanism does not care what they are.
    let rgb = synth_rgb();
    let pixels = (W * H) as u64;
    let raw_hash = blake3::hash(&rgb).into();

    // step 1: lossless PNG master — measured, unless the figures are carried.
    let (png_m, png_v) = match carried {
        Some(_) => (None, None),
        None => measure("encode.png", &rgb, |r| {
            let _ = std::hint::black_box(png_encode(r, H, W));
        }),
    };
    let png = png_encode(&rgb, H, W).expect("png");
    let png_hash: [u8; 32] = blake3::hash(&png).into();

    // step 2: published JPEG (from the master pixels) — measured, unless carried.
    let (jpeg_m, jpeg_v) = match carried {
        Some(_) => (None, None),
        None => measure("encode.jpeg", &rgb, |r| {
            let _ = std::hint::black_box(jpeg_encode(r, H, W, 90));
        }),
    };
    let jpeg = jpeg_encode(&rgb, H, W, 90).expect("jpeg");

    // Wrap the JPEG as a content-addressed wai.image.jpeg envelope.
    let manifest = Manifest::new("image", "replicate", ModelRequirement::new("wai.image.jpeg"), Conditioning::new("jpeg"))
        .with_target(serde_json::json!({ "name": "captured", "w": W, "h": H }));
    let envelope = Wai::new(manifest, jpeg.clone()).to_bytes().expect("wrap");
    let asset_hash = asset_content_hash(&envelope);

    // Each measured step carries its figure's class; a step whose
    // measurement failed is recorded unmetered, never with a made-up number.
    let step = |op: &str, input: [u8; 32], output: [u8; 32], m: Option<(u64, EnergyProvenance)>| {
        let s = ProvStep::new(op, input, output, pixels, m.as_ref().map_or(0, |(uj, _)| *uj));
        match m {
            Some((_, class)) => s.with_energy_class(class).expect("a usable figure is a valid label"),
            None => s,
        }
    };
    let mut steps = vec![
        ProvStep::new("origin.synthesize", [0u8; 32], raw_hash, pixels, 0),
        step("encode.png", raw_hash, png_hash, png_m.clone()),
        step("encode.jpeg", png_hash, asset_hash, jpeg_m.clone()),
    ];
    if let (Some(from), Some(path)) = (&carried, &carry) {
        steps = carry_figures(path, from, steps, asset_hash);
    }

    let signer = SigningKey::from_bytes(&[0x70u8; 32]); // 'p' for prov
    // A demo signer whose seed is published: a label, never origin-bound (jwp-receipts §7.1, §7.8).
    let signer_id = "did:example:wai-prov";
    let receipt = AssetReceipt::build(
        &signer,
        signer_id,
        asset_hash,
        "image",
        "wai.image.jpeg",
        0,
        steps,
        None,
    );
    assert!(receipt.verify(), "emitted receipt must verify");
    assert!(receipt.matches_asset(&envelope), "receipt must vouch for the asset");

    std::fs::write(dir.join("captured.jpeg.wai"), &envelope).expect("write asset");
    std::fs::write(dir.join("captured.receipt.json"), receipt.to_json()).expect("write receipt");

    println!("captured.jpeg.wai     : {} B envelope (wai.image.jpeg), asset {}…", envelope.len(), &hex(&asset_hash)[..16]);
    // The origin step is never metered, so the total covers the measured steps
    // only, and an all-unmetered lineage is UNMETERED, never 0 µJ.
    println!("captured.receipt.json : {} steps, total {}, receipt {}…", receipt.steps.len(), receipt.total_label(), &hex(&receipt.receipt_hash())[..16]);
    if let Some(path) = &carry {
        let shown = |s: &ProvStep| match (s.joules_micro, &s.energy_provenance) {
            (0, _) => "UNMETERED".to_string(),
            (uj, Some(c)) => format!("{uj} µJ ({})", c.label()),
            (uj, None) => format!("{uj} µJ (unlabelled)"),
        };
        println!("  encode.png  : {}   encode.jpeg : {}   (carried from {path})", shown(&receipt.steps[1]), shown(&receipt.steps[2]));
        println!("report            : not written — the figures are carried, and this run took no bracket for a report to restate");
        return;
    }
    let shown = |m: &Option<(u64, EnergyProvenance)>| match m {
        Some((uj, c)) => format!("{uj} µJ ({})", c.label()),
        None => "UNMETERED".to_string(),
    };
    println!("  encode.png  : {}   encode.jpeg : {}", shown(&png_m), shown(&jpeg_m));

    // One report per measured step, by its index: the first chained to the
    // previous run's report when `--previous-report` names one, each later one
    // to the report before it; a step recorded unmetered gets none, and says why.
    let mut chain = report_args.clone();
    for (index, values) in [(1u64, png_v), (2, jpeg_v)] {
        let path = dir.join(format!("captured.step{index}.report.json"));
        let path_s = path.to_string_lossy().into_owned();
        if meter_report::write(&chain, &receipt, Some(index), values, power_reader(), &signer, signer_id, None, Some(&path_s))
            .is_some()
        {
            chain.previous_report = Some(path_s);
        }
    }
}

/// Take `--carry <receipt.json>` out of `argv`. `--carry` with no file after it, or
/// given twice, is a usage error (exit status 2), never a measured run.
fn take_carry(argv: &mut Vec<String>) -> Option<String> {
    let mut found = None;
    while let Some(i) = argv.iter().position(|a| a == "--carry") {
        if found.is_some() || !argv.get(i + 1).is_some_and(|p| !p.starts_with("--")) {
            eprintln!("usage error: --carry takes one earlier receipt's file, once");
            std::process::exit(2);
        }
        found = Some(argv.remove(i + 1));
        argv.remove(i);
    }
    found
}

/// The receipt `--carry` names. It must parse and verify.
fn read_carried(path: &str) -> AssetReceipt {
    let json = std::fs::read_to_string(path).unwrap_or_else(|e| panic!("read {path}: {e}"));
    let r = AssetReceipt::from_json(&json).unwrap_or_else(|| panic!("{path}: not an asset-provenance receipt"));
    assert!(r.verify(), "{path}: the receipt to carry from does not verify");
    r
}

/// This run's steps, each with the figure and the label `from` states for it.
/// Refused (exit status 1) unless `from` is a receipt for this lineage: the same
/// asset, media, capability and group, no parent, and step for step the same
/// operation, input and output hashes and work.
fn carry_figures(path: &str, from: &AssetReceipt, steps: Vec<ProvStep>, asset_hash: [u8; 32]) -> Vec<ProvStep> {
    let same_asset = from.asset_hash == asset_hash
        && from.media == "image"
        && from.capability == "wai.image.jpeg"
        && from.group_id == 0
        && from.parent_receipt_hash.is_none();
    let same_steps = from.steps.len() == steps.len()
        && from.steps.iter().zip(&steps).all(|(c, s)| {
            c.op == s.op && c.input_hash == s.input_hash && c.output_hash == s.output_hash && c.work == s.work
        });
    if !same_asset || !same_steps {
        eprintln!("{path}: not a receipt for this lineage (its asset, or its steps' operations, hashes or work, differ from this run's); nothing carried");
        std::process::exit(1);
    }
    steps
        .into_iter()
        .zip(&from.steps)
        .map(|(s, c)| ProvStep { joules_micro: c.joules_micro, energy_provenance: c.energy_provenance.clone(), ..s })
        .collect()
}

/// Deterministic 512×512 RGB: radial rings over a diagonal gradient.
fn synth_rgb() -> Vec<u8> {
    let mut out = Vec::with_capacity((W * H * 3) as usize);
    let (cx, cy) = (W as f32 / 2.0, H as f32 / 2.0);
    for y in 0..H {
        for x in 0..W {
            let dx = x as f32 - cx;
            let dy = y as f32 - cy;
            let d = (dx * dx + dy * dy).sqrt();
            let ring = ((d * 0.09).sin() * 0.5 + 0.5) * 255.0;
            let gx = x as f32 / W as f32;
            let gy = y as f32 / H as f32;
            out.push((ring * (0.4 + 0.6 * gx)) as u8);
            out.push((ring * (0.4 + 0.6 * gy)) as u8);
            out.push((ring * (0.4 + 0.6 * (1.0 - gx))) as u8);
        }
    }
    out
}

/// One encode's figure and class, bracketed by the §3 protocol
/// (`energy_burst`): `Some((joules_micro, ModelBased))` for a usable
/// figure, `None` — the unmetered marker — otherwise; and the bracket's values
/// for the step's measurement report.
fn measure<F>(label: &str, rgb: &[u8], op: F) -> (Option<(u64, EnergyProvenance)>, Option<BracketValues>)
where
    F: Fn(&Vec<u8>) + Send + Sync + 'static,
{
    let Some(b) = measure_with_macmon(label, &rgb.to_vec(), op) else {
        eprintln!("macmon unavailable; recording {label} UNMETERED. Install `brew install macmon` (a non-sudo power reader).");
        return (None, None);
    };
    eprint!("{}", b.report(label));
    b.warn_if_unpublishable();
    // A reported power seals `ModelBased`, never `OnChipCounter` (§3.1).
    let m = b.figure(CounterFiltering::Unknown);
    if m.is_none() {
        eprintln!(
            "  {label}: no usable figure ({:.1}x noise: {}); recorded UNMETERED.",
            b.ratio(),
            b.unsealed_reason().unwrap_or("refused")
        );
    }
    (m, b.values())
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}
