//! The integer side of the determinism-cost record (`research/determinism-cost.md`):
//! decode `WIH1` payloads (`wai.neural.int_hyper`, integer-payloads §4) through the one
//! integer decode path a sink runs (`int_payload::decode_int`, over an unpinned held set),
//! and say what was decoded, or time it.
//!
//! ```text
//! wai_det_cost decode <jobs>   one job per line, `<set-dir>\t<payload>\t<out.rgb or ->`;
//!                              prints `<payload>\t<sha256 of the RGB8>\t<width>\t<height>`,
//!                              or `<payload>\tREFUSED\t<code>`, per job
//! wai_det_cost serve <jobs>    the timed child of tools/wai_det_energy.py (below)
//! ```
//!
//! A set directory holds the three members integer-payloads §9 names: `entropy.json`,
//! `g_s.whm` and `h_s.whm`. Options, before the subcommand:
//!
//! - `--kernels fast|reference`: the kernel policy (`int_kernels::Policy`); the fast kernels
//!   take their worker count from `WAI_INT_THREADS` (all cores when it is unset).
//! - `--gpu` (feature `ops_gpu`): every graph a decode runs goes to the first adapter found
//!   (`int_ops::gpu::with_gpu`); a graph the device does not take runs on the interpreter.
//!   `decode` then appends `\tgpu=<runs on the device>\tcpu=<runs on the interpreter>`.
//!   The adapter is the first GPU found (not a software rasteriser), or the first whose key
//!   contains `WAI_DET_GPU_KEY`.
//!
//! `serve` decodes every job once, prints `ready <jobs>`, and then reads commands from
//! standard input: `go <k>` decodes every job `k` times and prints `done <decodes>
//! <wall-ns>`; `exit` ends it. Between commands it blocks on its input, so every thread it
//! ran is off its core when the sampler reads the process's counters. Every decode's output
//! is hashed and compared with the first decode of its job: a time is never reported for a
//! different output (the process exits with status 3 instead).

use std::collections::BTreeMap;
use std::io::{BufRead, Write};
use std::path::Path;
use std::time::Instant;

use sha2::{Digest, Sha256};
use wai::int_kernels::{with_policy, Policy};
use wai::int_payload::{decode_int, resolve, IntDecoded, IntParams};
use wai::pin::HeldBytes;

const CAP: &str = "wai.neural.int_hyper";
const MEMBERS: [&str; 3] = ["entropy.json", "g_s.whm", "h_s.whm"];

struct Job {
    payload_path: String,
    out: Option<String>,
    payload: Vec<u8>,
    set: BTreeMap<String, Vec<u8>>,
}

fn read_jobs(list: &str) -> Vec<Job> {
    let text = std::fs::read_to_string(list).unwrap_or_else(|e| die(&format!("{list}: {e}")));
    let mut sets: BTreeMap<String, BTreeMap<String, Vec<u8>>> = BTreeMap::new();
    let mut jobs = Vec::new();
    for line in text.lines().filter(|l| !l.trim().is_empty()) {
        let f: Vec<&str> = line.split('\t').collect();
        if f.len() != 3 {
            die(&format!("a job line has {} fields, not 3: {line:?}", f.len()));
        }
        let set = sets
            .entry(f[0].to_string())
            .or_insert_with(|| {
                MEMBERS
                    .iter()
                    .map(|m| {
                        let p = Path::new(f[0]).join(m);
                        (m.to_string(), std::fs::read(&p).unwrap_or_else(|e| die(&format!("{}: {e}", p.display()))))
                    })
                    .collect()
            })
            .clone();
        let payload = std::fs::read(f[1]).unwrap_or_else(|e| die(&format!("{}: {e}", f[1])));
        jobs.push(Job { payload_path: f[1].to_string(), out: (f[2] != "-").then(|| f[2].to_string()), payload, set });
    }
    jobs
}

fn die(why: &str) -> ! {
    eprintln!("wai_det_cost: {why}");
    std::process::exit(2)
}

/// A job's set as the decode reads it.
fn view(job: &Job) -> BTreeMap<String, &[u8]> {
    job.set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect()
}

/// The parameters of an unpinned decode over `view`: resolved once per job, outside any
/// timed span, so a timed decode reads its set from the store as a sink's second decode
/// of a payload under one set does.
fn params<'a>(view: &'a BTreeMap<String, &'a [u8]>) -> IntParams<'a> {
    resolve(CAP, None, &[HeldBytes::Set { listing: b"", files: view }], None).unwrap_or_else(|_| IntParams::none())
}

/// One decode.
fn decode_one(job: &Job, params: &IntParams<'_>, policy: Policy) -> Decoded {
    match with_policy(policy, || decode_int(CAP, &job.payload, params)) {
        Ok(IntDecoded::Image { rgb, width, height }) => Ok((rgb, width, height)),
        Ok(_) => Err("not_an_image".into()),
        Err(e) => Err(e.code().to_string()),
    }
}

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

#[cfg(all(feature = "ops_gpu", not(target_arch = "wasm32")))]
mod device {
    use std::sync::Arc;
    use wai::int_ops::gpu::{self, wai_ops_gpu};

    pub struct Dev(Arc<wai_ops_gpu::Gpu>);

    pub fn open() -> Option<Dev> {
        let d = wai_ops_gpu::blocking::discover();
        for (a, why) in &d.refused {
            eprintln!("adapter refused: key={} {} reason={why:?}", a.key(), a.describe());
        }
        // The first GPU (never a software rasteriser, which a host can offer beside it), or
        // the first adapter whose key contains WAI_DET_GPU_KEY.
        let want = std::env::var("WAI_DET_GPU_KEY").ok();
        let g = d
            .found
            .into_iter()
            .find(|g| match &want {
                Some(k) => g.adapter().key().contains(k.as_str()),
                None => g.adapter().is_gpu(),
            })
            .unwrap_or_else(|| super::die("--gpu: no adapter was found"));
        eprintln!("adapter: key={} {}", g.adapter().key(), g.adapter().describe());
        Some(Dev(Arc::new(g)))
    }

    /// `f` with every graph on the device: its result, and the graph runs on the device
    /// and on the interpreter.
    pub fn run<R>(d: &Dev, f: impl FnOnce() -> R) -> (R, u64, u64) {
        let (r, s) = gpu::with_gpu(&d.0, f);
        if !s.errors.is_empty() {
            super::die(&format!("device errors: {:?}", s.errors));
        }
        (r, s.on_gpu(), s.on_cpu())
    }
}

/// No device without `ops_gpu` on a native target: `--gpu` is refused there.
#[cfg(not(all(feature = "ops_gpu", not(target_arch = "wasm32"))))]
mod device {
    pub type Dev = std::convert::Infallible;
    pub fn open() -> Option<Dev> {
        super::die("--gpu needs the ops_gpu feature on a native target")
    }
    pub fn run<R>(d: &Dev, _: impl FnOnce() -> R) -> (R, u64, u64) {
        match *d {}
    }
}

/// One decode's outcome: the RGB8 and its size, or the refusal's code.
type Decoded = Result<(Vec<u8>, u32, u32), String>;

fn main() {
    let mut args: Vec<String> = std::env::args().skip(1).collect();
    let mut policy = Policy::default();
    let mut dev: Option<device::Dev> = None;
    while let Some(a) = args.first().cloned() {
        match a.as_str() {
            "--kernels" => {
                policy = match args.get(1).map(String::as_str) {
                    Some("fast") => Policy::default(),
                    Some("reference") => Policy::reference(),
                    other => die(&format!("--kernels takes fast or reference, not {other:?}")),
                };
                args.drain(..2);
            }
            "--gpu" => {
                dev = device::open();
                args.remove(0);
            }
            _ => break,
        }
    }
    let (cmd, list) = match args.as_slice() {
        [c, l] => (c.as_str(), l.as_str()),
        _ => die("usage: wai_det_cost [--kernels fast|reference] [--gpu] decode|serve <jobs>"),
    };
    let jobs = read_jobs(list);
    let views: Vec<BTreeMap<String, &[u8]>> = jobs.iter().map(view).collect();
    let prms: Vec<IntParams<'_>> = views.iter().map(params).collect();
    let run = |i: usize| -> (Decoded, u64, u64) {
        match &dev {
            Some(d) => device::run(d, || decode_one(&jobs[i], &prms[i], policy)),
            None => (decode_one(&jobs[i], &prms[i], policy), 0, 0),
        }
    };
    let stdout = std::io::stdout();
    match cmd {
        "decode" => {
            for (i, job) in jobs.iter().enumerate() {
                let (r, on_gpu, on_cpu) = run(i);
                let mut line = match r {
                    Ok((rgb, w, h)) => {
                        if let Some(o) = &job.out {
                            std::fs::write(o, &rgb).unwrap_or_else(|e| die(&format!("{o}: {e}")));
                        }
                        format!("{}\t{}\t{w}\t{h}", job.payload_path, sha256_hex(&rgb))
                    }
                    Err(code) => format!("{}\tREFUSED\t{code}", job.payload_path),
                };
                if dev.is_some() {
                    line.push_str(&format!("\tgpu={on_gpu}\tcpu={on_cpu}"));
                }
                let mut o = stdout.lock();
                writeln!(o, "{line}").unwrap();
                o.flush().unwrap();
            }
        }
        "serve" => {
            let want: Vec<String> = jobs
                .iter()
                .enumerate()
                .map(|(i, j)| match run(i).0 {
                    Ok((rgb, ..)) => sha256_hex(&rgb),
                    Err(code) => die(&format!("{}: refused ({code}); nothing to time", j.payload_path)),
                })
                .collect();
            println!("ready {}", jobs.len());
            std::io::stdout().flush().unwrap();
            let stdin = std::io::stdin();
            for line in stdin.lock().lines() {
                let line = line.unwrap_or_default();
                let mut w = line.split_whitespace();
                match (w.next(), w.next()) {
                    (Some("go"), Some(k)) => {
                        let k: u64 = k.parse().unwrap_or_else(|_| die(&format!("go takes a count, not {k:?}")));
                        let t0 = Instant::now();
                        let mut n = 0u64;
                        for _ in 0..k {
                            for (i, (j, h)) in jobs.iter().zip(&want).enumerate() {
                                let got = run(i).0.map(|(rgb, ..)| sha256_hex(&rgb));
                                if got.as_ref() != Ok(h) {
                                    println!("MISMATCH {} {got:?} {h}", j.payload_path);
                                    std::process::exit(3);
                                }
                                n += 1;
                            }
                        }
                        let ns = t0.elapsed().as_nanos();
                        println!("done {n} {ns}");
                        std::io::stdout().flush().unwrap();
                    }
                    (Some("exit"), _) | (None, _) => return,
                    _ => die(&format!("unknown command {line:?}")),
                }
            }
        }
        other => die(&format!("unknown subcommand {other:?}")),
    }
}
