//! The reference deriving edge (`extensions/deriving-edge.md`).
//!
//!   wai_edge derive --bundle <in.json> --out <out.json> --key <seed file> --signer-id <id>
//!                   --derivation <id> [--params <json>] [--encoder <json>] --output <capability>
//!                   --input <content hash>[@<frame>][=<capability>] ...
//!                   --track <n> --group <n> --time <date-time> --source <key hex> ...
//!                   [--head <origin>=<file>] ... [--energy <json>] [--origin <label>]
//!                   [--personal <digest>] ... [--synthesis <content hash>=<label>] ...
//!                   [--own-synthesis <label>]
//!
//!       Derive an object from objects of a bundle, publish it under the edge's key in a group
//!       of its own track, sign its record, and write the bundle with the output, its receipts
//!       and its record added. Each input is verified first (§3.3): a source object's receipts
//!       must verify and be signed by a `--source` key, and a derived input's chain must verify
//!       by §7. The group's parent is the edge's latest group on the track in the bundle. The
//!       edge measures nothing: its energy is the unmetered marker unless `--energy` gives a
//!       figure a meter took, as a record's `energy` member.
//!
//!       Every provenance item of every input is carried forward (§6.1): the bundle's claims
//!       about it and the C2PA Manifest Stores bound to it, in the system compartment unless
//!       `--personal` names the item's digest; the stores embedded in it; and every item its
//!       record carries. `--synthesis` states the AI-origin label a source input's provenance
//!       gives (§6.2), and `--own-synthesis` the label the derivation's own step warrants.
//!
//!   wai_edge verify --bundle <b.json> --subject <content hash> --source <key hex> ...
//!                   [--head <origin>=<file>] ... [--refuse-attested] [--require-recompute]
//!                   [--require-authority]
//!
//!       Verify the object `--subject` back to its source (§7) and print the outcome as JSON.
//!       Exit status 0 when it is accepted, 1 when it is refused.
//!
//! A bundle is a JSON object: `records`, record texts; `receipts`, group and object receipts;
//! `objects`, envelopes by content hash, in lowercase hexadecimal; and, where there are any,
//! `claims`, claim texts, and `c2pa`, C2PA Manifest Stores in hexadecimal by the content hash
//! of the object each is bound to. The key file holds the 32-byte Ed25519 seed as 64
//! hexadecimal digits.

use std::collections::BTreeMap;
use std::process::ExitCode;

use ed25519_dalek::SigningKey;
use serde_json::{json, Map, Value};
use wai::derive_edge::edge::{self, InputObject, Publish};
use wai::derive_edge::verify::{self, Held, Policy};
use wai::derive_edge::{claim_about, sha256, Compartment, Derivation, Encoder, Energy, Params, ProvKind, Record, Synthesis};
use wai::jwp_delivery::{GroupReceipt, ObjectReceipt};
use wai::jwp_hash::content_hash;

fn usage() -> ExitCode {
    eprintln!("usage: wai_edge derive --bundle <in.json> --out <out.json> --key <seed file> --signer-id <id> --derivation <id> [--params <json>] [--encoder <json>] --output <capability> --input <hash>[@<frame>][=<capability>] ... --track <n> --group <n> --time <date-time> --source <key> ... [--head <origin>=<file>] ... [--energy <json>] [--origin <label>]");
    eprintln!("       wai_edge verify --bundle <b.json> --subject <hash> --source <key> ... [--head <origin>=<file>] ... [--refuse-attested] [--require-recompute] [--require-authority]");
    ExitCode::from(2)
}

fn die(msg: impl std::fmt::Display) -> ExitCode {
    eprintln!("error: {msg}");
    ExitCode::from(2)
}

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

fn unhex(s: &str) -> Option<Vec<u8>> {
    if !s.len().is_multiple_of(2) || !s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) {
        return None;
    }
    (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).ok()).collect()
}

fn hash32(s: &str) -> Result<[u8; 32], String> {
    unhex(s).and_then(|b| b.try_into().ok()).ok_or_else(|| format!("{s}: not 64 lowercase hexadecimal digits"))
}

/// The options: each name with every value given for it, and the flags.
struct Opts(BTreeMap<String, Vec<String>>);

impl Opts {
    fn parse(args: &[String]) -> Result<Opts, String> {
        let mut m: BTreeMap<String, Vec<String>> = BTreeMap::new();
        let mut i = 0;
        while i < args.len() {
            let a = args[i].strip_prefix("--").ok_or_else(|| format!("{}: not an option", args[i]))?;
            if a.starts_with("refuse-") || a.starts_with("require-") {
                m.entry(a.to_owned()).or_default().push(String::new());
                i += 1;
                continue;
            }
            let v = args.get(i + 1).ok_or_else(|| format!("--{a} needs a value"))?;
            m.entry(a.to_owned()).or_default().push(v.clone());
            i += 2;
        }
        Ok(Opts(m))
    }

    fn one(&self, k: &str) -> Result<&str, String> {
        match self.0.get(k).map(Vec::as_slice) {
            Some([v]) => Ok(v),
            Some(_) => Err(format!("--{k} given more than once")),
            None => Err(format!("--{k} is required")),
        }
    }

    fn opt(&self, k: &str) -> Option<&str> {
        self.0.get(k).and_then(|v| v.first()).map(String::as_str)
    }

    fn all(&self, k: &str) -> &[String] {
        self.0.get(k).map_or(&[], Vec::as_slice)
    }

    fn flag(&self, k: &str) -> bool {
        self.0.contains_key(k)
    }
}

/// A bundle of records, receipts and objects.
struct Bundle {
    records: Vec<String>,
    receipts: Vec<Value>,
    objects: BTreeMap<[u8; 32], Vec<u8>>,
    claims: Vec<String>,
    c2pa: BTreeMap<[u8; 32], Vec<Vec<u8>>>,
}

impl Bundle {
    fn read(path: &str) -> Result<Bundle, String> {
        let text = std::fs::read_to_string(path).map_err(|e| format!("{path}: {e}"))?;
        let v: Value = serde_json::from_str(&text).map_err(|e| format!("{path}: {e}"))?;
        let records = v["records"]
            .as_array()
            .ok_or("no records")?
            .iter()
            .map(|r| r.as_str().map(String::from).ok_or("a record that is not a string"))
            .collect::<Result<_, _>>()?;
        let receipts = v["receipts"].as_array().ok_or("no receipts")?.clone();
        let mut objects = BTreeMap::new();
        for (h, b) in v["objects"].as_object().ok_or("no objects")? {
            objects.insert(hash32(h)?, unhex(b.as_str().unwrap_or("")).ok_or("an object that is not hex")?);
        }
        let claims = match v.get("claims") {
            None => Vec::new(),
            Some(c) => c
                .as_array()
                .ok_or("claims is not an array")?
                .iter()
                .map(|t| t.as_str().map(String::from).ok_or("a claim that is not a string"))
                .collect::<Result<_, _>>()?,
        };
        let mut c2pa = BTreeMap::new();
        if let Some(m) = v.get("c2pa") {
            for (h, stores) in m.as_object().ok_or("c2pa is not an object")? {
                let stores = stores
                    .as_array()
                    .ok_or("a c2pa entry is not an array")?
                    .iter()
                    .map(|b| unhex(b.as_str().unwrap_or("")).ok_or("a store that is not hex"))
                    .collect::<Result<Vec<_>, _>>()?;
                c2pa.insert(hash32(h)?, stores);
            }
        }
        Ok(Bundle { records, receipts, objects, claims, c2pa })
    }

    fn write(&self, path: &str) -> Result<(), String> {
        let v = json!({
            "records": self.records,
            "receipts": self.receipts,
            "objects": self.objects.iter().map(|(h, b)| (hex(h), json!(hex(b)))).collect::<Map<String, Value>>(),
            "claims": self.claims,
            "c2pa": self.c2pa.iter().map(|(h, st)| (hex(h), json!(st.iter().map(|b| hex(b)).collect::<Vec<_>>()))).collect::<Map<String, Value>>(),
        });
        std::fs::write(path, serde_json::to_string_pretty(&v).expect("JSON") + "\n").map_err(|e| format!("{path}: {e}"))
    }

    fn groups(&self) -> Vec<GroupReceipt> {
        self.receipts
            .iter()
            .filter(|r| r["kind"] == "group")
            .filter_map(|r| GroupReceipt::from_json(r).ok())
            .collect()
    }

    /// The object receipt naming `hash` and its group, both verifying.
    fn receipts_of(&self, hash: &[u8; 32]) -> Option<(GroupReceipt, ObjectReceipt)> {
        let groups = self.groups();
        self.receipts
            .iter()
            .filter(|r| r["kind"] == "object")
            .filter_map(|r| ObjectReceipt::from_json(r).ok())
            .filter(|o| o.content_hash == *hash)
            .find_map(|o| {
                let g = groups.iter().find(|g| {
                    g.track_alias == o.track_alias && g.group_id == o.group_id && g.verify().is_ok() && o.verify_inclusion(g).is_ok()
                })?;
                Some((g.clone(), o))
            })
    }

    /// The position of the last record whose output is `hash`.
    fn record_of(&self, hash: &[u8; 32]) -> Option<usize> {
        self.records.iter().rposition(|t| Record::read(t).is_ok_and(|r| r.output.content_hash == *hash))
    }

    fn held(&self, sources: &[[u8; 32]], heads: &BTreeMap<String, String>, policy: Policy) -> Held {
        Held {
            records: self.records.clone(),
            receipts: self.receipts.clone(),
            objects: self.objects.clone(),
            sources: sources.to_vec(),
            heads: heads.clone(),
            policy,
            claims: self.claims.clone(),
            c2pa: self.c2pa.clone(),
        }
    }
}

/// The source keys and the key-log heads by origin.
type SourcesAndHeads = (Vec<[u8; 32]>, BTreeMap<String, String>);

fn sources_and_heads(o: &Opts) -> Result<SourcesAndHeads, String> {
    let sources = o.all("source").iter().map(|s| hash32(s)).collect::<Result<Vec<_>, _>>()?;
    let mut heads = BTreeMap::new();
    for h in o.all("head") {
        let (origin, file) = h.split_once('=').ok_or_else(|| format!("--head {h}: not <origin>=<file>"))?;
        heads.insert(origin.to_owned(), std::fs::read_to_string(file).map_err(|e| format!("{file}: {e}"))?);
    }
    Ok((sources, heads))
}

fn json_arg(o: &Opts, k: &str) -> Result<Option<Value>, String> {
    o.opt(k).map(|t| serde_json::from_str(t).map_err(|e| format!("--{k}: {e}"))).transpose()
}

fn derive(o: &Opts) -> Result<String, String> {
    let mut bundle = Bundle::read(o.one("bundle")?)?;
    let seed = std::fs::read_to_string(o.one("key")?).map_err(|e| format!("--key: {e}"))?;
    let key = SigningKey::from_bytes(&hash32(seed.trim())?);
    let signer_id = o.one("signer-id")?;
    let d = Derivation::from_id(o.one("derivation")?).ok_or("--derivation: not a derivation of deriving-edge §4")?;
    let (sources, heads) = sources_and_heads(o)?;

    // The inputs, each verified before the edge derives from it (§3.3).
    struct In {
        hash: [u8; 32],
        frame: Option<u64>,
        capability: Option<String>,
        group: GroupReceipt,
        object: ObjectReceipt,
        record: Option<Record>,
    }
    let personal = o.all("personal").iter().map(|h| hash32(h)).collect::<Result<Vec<_>, _>>()?;
    let compartment = |d: &[u8; 32]| if personal.contains(d) { Compartment::Personal } else { Compartment::System };
    let mut labels = BTreeMap::new();
    for l in o.all("synthesis") {
        let (h, label) = l.split_once('=').ok_or_else(|| format!("--synthesis {l}: not <content hash>=<label>"))?;
        labels.insert(hash32(h)?, Synthesis::from_label(label).ok_or_else(|| format!("--synthesis {l}: not a label"))?);
    }
    let own = o.opt("own-synthesis").map(|l| Synthesis::from_label(l).ok_or("--own-synthesis: not a label")).transpose()?;
    let mut ins = Vec::new();
    for spec in o.all("input") {
        let (spec, capability) = match spec.split_once('=') {
            Some((s, c)) => (s, Some(c.to_owned())),
            None => (spec.as_str(), None),
        };
        let (h, frame) = match spec.split_once('@') {
            Some((h, f)) => (h, Some(f.parse::<u64>().map_err(|_| format!("--input {spec}: a frame is a number"))?)),
            None => (spec, None),
        };
        let hash = hash32(h)?;
        let env = bundle.objects.get(&hash).ok_or_else(|| format!("input {h}: not in the bundle"))?;
        if content_hash(env) != hash {
            return Err(format!("input {h}: its envelope does not hash to its name"));
        }
        let (group, object) = bundle.receipts_of(&hash).ok_or_else(|| format!("input {h}: no receipts that verify"))?;
        let record = match bundle.record_of(&hash) {
            None => {
                if !sources.contains(&group.signer_pubkey) {
                    return Err(format!("input {h}: signed by {}, which is not a --source", hex(&group.signer_pubkey)));
                }
                None
            }
            Some(r) => {
                verify::verify(&bundle.held(&sources, &heads, Policy::default()), r)
                    .map_err(|f| format!("input {h}: its chain is refused: {} at {}", f.reason, f.at.to_value()))?;
                Some(Record::read(&bundle.records[r]).expect("a verified record reads"))
            }
        };
        ins.push(In { hash, frame: frame.or(d.evaluates().then_some(0)), capability, group, object, record });
    }
    if ins.is_empty() {
        return Err("no --input".into());
    }
    let caps: Vec<String> = ins
        .iter()
        .map(|i| {
            i.capability.clone().map_or_else(
                || wai::derive_edge::ops::wai1_capability(&bundle.objects[&i.hash]).ok_or_else(|| "an input that is not a WAI1; name its capability".to_string()),
                Ok,
            )
        })
        .collect::<Result<_, _>>()?;
    let params = Params::from_value(d, &json_arg(o, "params")?.unwrap_or(json!({})), ins.len()).map_err(|r| format!("--params: {r}"))?;
    let encoder = json_arg(o, "encoder")?.map(|v| Encoder::from_value(&v)).transpose().map_err(|r| format!("--encoder: {r}"))?;
    let energy = match json_arg(o, "energy")? {
        None => None,
        Some(v) => Energy::from_value(&v).map_err(|r| format!("--energy: {r}"))?,
    };
    let track: u64 = o.one("track")?.parse().map_err(|_| "--track is a number")?;
    let group: u64 = o.one("group")?.parse().map_err(|_| "--group is a number")?;
    let me = key.verifying_key().to_bytes();
    let parent = bundle
        .groups()
        .into_iter()
        .filter(|g| g.signer_pubkey == me && g.track_alias == track && g.group_id < group)
        .max_by_key(|g| g.group_id)
        .map(|g| g.link_hash());
    // The provenance each input came with beside its envelope (§6.1).
    type Items = Vec<(ProvKind, [u8; 32], Compartment)>;
    let carried = |h: &[u8; 32]| -> Items {
        let claims = bundle.claims.iter().filter_map(|t| claim_about(t)).filter(|(s, _)| s == h).map(|(_, l)| (ProvKind::Claim, l));
        let stores = bundle.c2pa.get(h).into_iter().flatten().map(|b| (ProvKind::C2pa, sha256(b)));
        claims.chain(stores).map(|(k, d)| (k, d, compartment(&d))).collect()
    };
    let inputs: Vec<InputObject> = ins
        .iter()
        .zip(&caps)
        .map(|(i, c)| InputObject {
            envelope: &bundle.objects[&i.hash],
            capability: c,
            frame: i.frame,
            group: &i.group,
            object: &i.object,
            record: i.record.as_ref(),
            provenance: carried(&i.hash),
            synthesis: labels.get(&i.hash).copied(),
        })
        .collect();
    let p = Publish {
        key: &key,
        signer_id,
        track_alias: track,
        group_id: group,
        time: o.one("time")?,
        parent,
        origin: o.opt("origin").unwrap_or("edge"),
    };
    let derived = edge::derive(d, params, encoder, o.one("output")?, &inputs, energy, own, &p)?;
    let out_hash = derived.record.output.content_hash;
    let record_hash = derived.record.hash();
    bundle.records.push(derived.record.to_json());
    bundle.receipts.push(derived.group.to_json());
    bundle.receipts.push(derived.object.to_json());
    bundle.objects.insert(out_hash, derived.envelope);
    bundle.write(o.one("out")?)?;
    Ok(json!({ "output": hex(&out_hash), "record": hex(&record_hash) }).to_string())
}

fn verify_cmd(o: &Opts) -> Result<(String, bool), String> {
    let bundle = Bundle::read(o.one("bundle")?)?;
    let subject = hash32(o.one("subject")?)?;
    let (sources, heads) = sources_and_heads(o)?;
    let policy = Policy {
        refuse_attested: o.flag("refuse-attested"),
        require_recompute: o.flag("require-recompute"),
        require_authority: o.flag("require-authority"),
    };
    let r = bundle.record_of(&subject).ok_or("no record in the bundle derives the subject")?;
    Ok(match verify::verify(&bundle.held(&sources, &heads, policy), r) {
        Ok(f) => (json!({ "outcome": "accept", "facts": f.to_value() }).to_string(), true),
        Err(f) => (json!({ "outcome": "refuse", "reason": f.reason.code(), "at": f.at.to_value() }).to_string(), false),
    })
}

fn main() -> ExitCode {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let Some(cmd) = args.first() else { return usage() };
    let opts = match Opts::parse(&args[1..]) {
        Ok(o) => o,
        Err(e) => return die(e),
    };
    match cmd.as_str() {
        "derive" => match derive(&opts) {
            Ok(out) => {
                println!("{out}");
                ExitCode::SUCCESS
            }
            Err(e) => die(e),
        },
        "verify" => match verify_cmd(&opts) {
            Ok((out, accepted)) => {
                println!("{out}");
                if accepted { ExitCode::SUCCESS } else { ExitCode::from(1) }
            }
            Err(e) => die(e),
        },
        _ => usage(),
    }
}
