//! `wai encode avc|hevc`, `wai encoders` and `wai encode-record verify`: AVC and HEVC
//! through the encoders a device already has, and the signed record of who made the
//! bytes (extensions/platform-encode.md). A build with the feature `platform_encode`.

use std::path::{Path, PathBuf};
use wai::platform_encode::{self as pe, Api, Codec, Options, record};
use wai::{Conditioning, Manifest, ModelRequirement, Wai};

fn opt<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1).map(String::as_str))
}

fn read(p: &Path) -> Result<Vec<u8>, String> {
    std::fs::read(p).map_err(|e| format!("{}: {e}", p.display()))
}

fn write(p: &str, b: &[u8]) -> Result<(), String> {
    crate::out::write(p, b).map_err(|e| format!("{p}: {e}"))
}

fn options(args: &[String]) -> Result<Options, String> {
    let hardware = match opt(args, "--hardware") {
        None => None,
        Some("yes") => Some(true),
        Some("no") => Some(false),
        Some(h) => return Err(format!("--hardware {h}: yes or no")),
    };
    Ok(Options { library: opt(args, "--library").map(PathBuf::from), device: opt(args, "--device").map(String::from), hardware })
}

fn size(s: &str) -> Result<(u32, u32), String> {
    let (w, h) = s.split_once('x').ok_or_else(|| format!("--size {s}: WxH"))?;
    Ok((w.parse().map_err(|_| format!("--size {s}"))?, h.parse().map_err(|_| format!("--size {s}"))?))
}

/// The interfaces `--encoder auto` tries, in order: the platform's, a GPU driver's, the
/// video acceleration API, and the software library when a path is given.
fn candidates(choice: &str, opts: &Options) -> Result<Vec<Api>, String> {
    if choice == "auto" {
        let mut v = vec![Api::Platform, Api::GpuDriver, Api::Va];
        if opts.library.is_some() {
            v.push(Api::Software);
        }
        return Ok(v);
    }
    Api::from_name(choice).map(|a| vec![a]).ok_or_else(|| format!("--encoder {choice}: auto, platform, va, gpu-driver or software"))
}

pub fn usage() {
    eprintln!("wai encode avc|hevc <in.y4m> [-o out.wai] [--encoder auto|platform|va|gpu-driver|software]");
    eprintln!("            [--library <path>] [--device <node|ordinal>] [--hardware yes|no] [--kbps <n>] [--size WxH]");
    eprintln!("            [--idr-every <frames> | --idr-plan <plan.txt>] [--es <out.h264|.h265>]");
    eprintln!("            [--record <out.json> --key <seed>] [--meter <counter> [--rounds <odd n>]]");
    eprintln!("            AVC or HEVC through an encoder the device has, its stream checked, and a signed");
    eprintln!("            record naming the encoder, its version, the device and the driver (a build with");
    eprintln!("            `platform_encode`). Counters: task (macOS), rapl, gpu-board, rapl-and-gpu-board (Linux)");
    eprintln!("wai encoders [--library <path>]     what each encoder interface offers here");
    eprintln!("wai encode-record verify <record.json> [--stream <es>] [--source <in.y4m>]");
}

pub fn encode(args: &[String]) -> Result<(), String> {
    let codec = match args.first().map(String::as_str) {
        Some("avc") => Codec::Avc,
        Some("hevc") => Codec::Hevc,
        _ => return Err("encode avc|hevc <in.y4m>".into()),
    };
    let input = args.get(1).ok_or("encode avc|hevc <in.y4m>")?;
    let rest = &args[2..];
    let out = opt(rest, "-o").unwrap_or("out.wai");
    let kbps: u32 = opt(rest, "--kbps").map(|k| k.parse().map_err(|_| format!("--kbps {k}"))).transpose()?.unwrap_or(match codec {
        Codec::Avc => 6000,
        Codec::Hevc => 4000,
    });
    let opts = options(rest)?;
    let prepared = pe::prepare(&read(Path::new(input))?, codec, opt(rest, "--size").map(size).transpose()?, kbps).map_err(|e| e.to_string())?;
    let frames = prepared.source.frames;
    let idr = match (opt(rest, "--idr-plan"), opt(rest, "--idr-every")) {
        (Some(_), Some(_)) => return Err("--idr-plan or --idr-every, not both".into()),
        (Some(p), None) => pe::read_plan(&String::from_utf8_lossy(&read(Path::new(p))?)).map_err(|e| e.to_string())?,
        (None, e) => {
            let every = e.map(|e| e.parse::<u64>().map_err(|_| format!("--idr-every {e}"))).transpose()?;
            // A key frame every two seconds by default.
            pe::idr_every(frames, every.unwrap_or_else(|| (2 * u64::from(prepared.cfg.fps_num)).div_ceil(u64::from(prepared.cfg.fps_den)).max(1)))
        }
    };
    let mut why = Vec::new();
    let mut chosen = None;
    for api in candidates(opt(rest, "--encoder").unwrap_or("auto"), &opts)? {
        let run = || pe::encode(api, &prepared.cfg, &opts, prepared.frames.iter().map(Vec::as_slice), &idr);
        let r = match opt(rest, "--meter") {
            None => run().map(|e| (e, None)).map_err(|e| e.to_string()),
            Some(c) => {
                let rounds: usize = opt(rest, "--rounds").map(|r| r.parse().map_err(|_| format!("--rounds {r}"))).transpose()?.unwrap_or(5);
                let mut counter = pe::energy::open(c)?;
                // The figure is of one encode of this input with these settings; an
                // encoder that is not a pure function of its input may make other bytes
                // in each round, and the record states the last round's stream.
                let measured = pe::energy::measure(counter.as_mut(), rounds, std::time::Duration::from_millis(1500), || run().map_err(|e| e.to_string()));
                measured.map(|(energy, rounds, e)| {
                    for (i, r) in rounds.iter().enumerate() {
                        eprintln!("round {i}: {:?} µJ over {} µs (idle window {} µs)", r.uj, r.busy_us, r.idle_us);
                    }
                    (e, energy)
                })
            }
        };
        match r {
            Ok(x) => {
                chosen = Some(x);
                break;
            }
            Err(e) => why.push(format!("{}: {e}", api.name())),
        }
    }
    let (enc, energy) = chosen.ok_or_else(|| format!("no encoder made the stream: {}", why.join("; ")))?;
    if let Some(es) = opt(rest, "--es") {
        write(es, &enc.stream)?;
    }
    let (w, h) = (prepared.cfg.width, prepared.cfg.height);
    let manifest = Manifest::new("video", "replicate", ModelRequirement::new(codec.capability()), Conditioning::new(codec.name())).with_target(serde_json::json!({"frames": frames, "h": h, "w": w}));
    let n = Wai::new(manifest, enc.stream.clone()).write(out).map_err(|e| e.to_string())?;
    let id = &enc.identity;
    println!("encoded → {out}  ({n} bytes, capability {}, {} access units)", codec.capability(), enc.units.len());
    println!("encoder  {} {} ({}), {}", id.api, id.encoder, id.version, if id.hardware { "hardware" } else { "software" });
    println!("device   {}", id.device);
    println!("driver   {}", id.driver);
    if let Some(l) = &id.library_sha256 {
        println!("library  sha256 {l}");
    }
    println!("stream   sha256 {}  {} bytes  {} µs", pe::sha256_hex(&enc.stream), enc.stream.len(), enc.wall_us);
    if let Some(e) = &energy {
        println!("energy   {} µJ ({}, {}, {}, counter {}, ±{} ppm over {} µs)", e.joules_micro, e.class, e.coverage, e.attribution, e.counter, e.uncertainty.map_or(0, |u| u.0), e.uncertainty.map_or(0, |u| u.1));
    }
    if let Some(path) = opt(rest, "--record") {
        let key = crate::dx::key::signing_key(opt(rest, "--key").ok_or("--record needs --key")?, true).map_err(|e| e.message)?;
        let issued = opt(rest, "--issued-at").map(String::from).unwrap_or_else(record::now_rfc3339);
        let body = record::body(&enc, &prepared.cfg, &prepared.source, &idr, energy.as_ref(), None, &issued);
        let signed = record::sign(body, &key);
        let text = record::to_text(&signed);
        record::verify(&text, &record::Held { stream: Some(&enc.stream), source: None }).map_err(|r| format!("the record written does not verify: {r}"))?;
        write(path, text.as_bytes())?;
        println!("record   {path}  digest {}", record::digest(&signed));
    }
    Ok(())
}

pub fn encoders(args: &[String]) -> Result<(), String> {
    for p in pe::probe(&options(args)?) {
        match p.result {
            Ok(id) => println!("{:<10} {:<4} {} {} ({}) on {}; {}", p.api.name(), p.codec.name(), if id.hardware { "hardware" } else { "software" }, id.encoder, id.version, id.device, id.driver),
            Err(e) => println!("{:<10} {:<4} none: {e}", p.api.name(), p.codec.name()),
        }
    }
    Ok(())
}

/// Exit status: 0 verified, 1 refused (the reason printed), 2 a usage or file error.
pub fn record_cmd(args: &[String]) -> Result<bool, String> {
    if args.first().map(String::as_str) != Some("verify") {
        return Err("encode-record verify <record.json> [--stream <es>] [--source <in.y4m>]".into());
    }
    let path = args.get(1).ok_or("encode-record verify <record.json>")?;
    let text = String::from_utf8(read(Path::new(path))?).map_err(|_| format!("{path}: not UTF-8"))?;
    let stream = opt(args, "--stream").map(|p| read(Path::new(p))).transpose()?;
    let src = match opt(args, "--source") {
        Some(p) => {
            let (hd, frames) = pe::y4m::read_all(&read(Path::new(p))?).map_err(|e| e.to_string())?;
            Some((hd, frames.concat()))
        }
        None => None,
    };
    match record::verify(&text, &record::Held { stream: stream.as_deref(), source: src.as_ref().map(|(h, p)| (h, p.as_slice())) }) {
        Ok(v) => {
            println!("verified: {} by {} {} ({}) on {}; digest {}", v["capability"].as_str().unwrap_or_default(), v["encoder"]["api"].as_str().unwrap_or_default(), v["encoder"]["name"].as_str().unwrap_or_default(), v["encoder"]["version"].as_str().unwrap_or_default(), v["encoder"]["device"].as_str().unwrap_or_default(), record::digest(&v));
            Ok(true)
        }
        Err(r) => {
            println!("refused: {r}");
            Ok(false)
        }
    }
}
