//! Verification of a derived object back to its source (`extensions/deriving-edge.md` §7,
//! §8): each hop of its chain read, its signature checked, its output bound to the edge's
//! delivery receipts, its inputs to theirs and, hop by hop, to the records that derived
//! them; each `replicate` hop whose inputs are held re-executed and compared byte for byte;
//! the edge's authority from its sources' origins; and the sink's policy. The facts of an
//! accepted chain are §7.4's.

use super::ops::{self, Samples};
use super::{
    claim_about, held_hash, sha256, Check, Derivation, Params, ProvItem, ProvKind, Reason, Record, Synthesis, AUTHORITY,
    MAX_HOPS,
};
use crate::claim_receipt::hex;
use crate::jwp_delivery::{GroupReceipt, ObjectLabel, ObjectReceipt};
use crate::jwp_hash::content_hash;
use crate::key_log::{EntryStatus, Revision};
use crate::signer_id::{classify, SignerIdForm};
use serde_json::{json, Value};
use std::collections::{BTreeMap, BTreeSet};

/// The sink's policy (§7.1).
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Policy {
    /// `attested: refuse`.
    pub refuse_attested: bool,
    /// `recompute: required`.
    pub require_recompute: bool,
    /// `authority: required`.
    pub require_authority: bool,
}

impl Policy {
    /// Its JSON form, every member written.
    pub fn to_value(&self) -> Value {
        json!({
            "attested": if self.refuse_attested { "refuse" } else { "accept" },
            "recompute": if self.require_recompute { "required" } else { "when-held" },
            "authority": if self.require_authority { "required" } else { "when-given" },
        })
    }

    /// Read its JSON form: an absent member takes its default.
    pub fn from_value(v: &Value) -> Option<Policy> {
        let m = v.as_object()?;
        let pick = |k: &str, default: &str, other: &str| -> Option<bool> {
            match m.get(k).map(|v| v.as_str()) {
                None => Some(false),
                Some(Some(s)) if s == default => Some(false),
                Some(Some(s)) if s == other => Some(true),
                _ => None,
            }
        };
        Some(Policy {
            refuse_attested: pick("attested", "accept", "refuse")?,
            require_recompute: pick("recompute", "when-held", "required")?,
            require_authority: pick("authority", "when-given", "required")?,
        })
    }
}

/// What a verifier holds (§7.1).
#[derive(Clone, Debug, Default)]
pub struct Held {
    pub records: Vec<String>,
    pub receipts: Vec<Value>,
    pub objects: BTreeMap<[u8; 32], Vec<u8>>,
    pub sources: Vec<[u8; 32]>,
    pub heads: BTreeMap<String, String>,
    pub policy: Policy,
    /// Claim texts it holds (jwp-receipts §2.4.3): those about an input are its provenance.
    pub claims: Vec<String>,
    /// C2PA Manifest Stores it holds bound to an object, by the object's content hash.
    pub c2pa: BTreeMap<[u8; 32], Vec<Vec<u8>>>,
}

/// Where a refusal happened (§7.2): a record by its position among the records given.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Place {
    Record(usize),
    Output(usize),
    Input(usize, usize),
    Recompute(usize),
    Authority(usize, Option<String>),
    Policy(usize),
}

impl Place {
    pub fn to_value(&self) -> Value {
        match self {
            Place::Record(r) => json!(["record", r]),
            Place::Output(r) => json!(["record", r, "output"]),
            Place::Input(r, i) => json!(["record", r, "input", i]),
            Place::Recompute(r) => json!(["record", r, "recompute"]),
            Place::Authority(r, o) => json!(["record", r, "authority", o]),
            Place::Policy(r) => json!(["record", r, "policy"]),
        }
    }
}

/// A refusal: why, and where.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Refusal {
    pub reason: Reason,
    pub at: Place,
}

fn refuse<T>(reason: Reason, at: Place) -> Result<T, Refusal> {
    Err(Refusal { reason, at })
}

/// How a hop was held to its record (§7.4).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Checked {
    Recomputed,
    NotRecomputed,
    Attested,
}

impl Checked {
    pub fn label(self) -> &'static str {
        match self {
            Checked::Recomputed => "recomputed",
            Checked::NotRecomputed => "not-recomputed",
            Checked::Attested => "attested",
        }
    }
}

/// A hop's standing under §8.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Authority {
    Confirmed,
    Unconfirmed,
    None,
}

impl Authority {
    pub fn label(self) -> &'static str {
        match self {
            Authority::Confirmed => "confirmed",
            Authority::Unconfirmed => "unconfirmed",
            Authority::None => "none",
        }
    }
}

/// One hop of an accepted chain (§7.4).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Hop {
    pub record: [u8; 32],
    pub signer: [u8; 32],
    pub derivation: Derivation,
    pub check: Check,
    pub checked: Checked,
    /// The figure and its class's label.
    pub energy: Option<(u64, String)>,
    pub authority: Authority,
    pub synthesis: Option<Synthesis>,
}

/// A source object an accepted chain reaches (§7.4).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Source {
    pub object: [u8; 32],
    pub signer: [u8; 32],
    pub signer_id: String,
}

/// What an accepted chain states (§7.4).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Facts {
    pub subject: [u8; 32],
    pub hops: Vec<Hop>,
    pub sources: Vec<Source>,
    /// Per class, strongest first: the class's label and the sum of its figures.
    pub energy: Vec<(String, u64)>,
    pub unmetered: u64,
    /// The provenance items the subject's record carries (§6.1).
    pub provenance: Vec<ProvItem>,
}

impl Facts {
    pub fn to_value(&self) -> Value {
        json!({
            "subject": hex(&self.subject),
            "hops": self.hops.iter().map(|h| json!({
                "record": hex(&h.record),
                "signer": hex(&h.signer),
                "derivation": h.derivation.id(),
                "check": h.check.label(),
                "checked": h.checked.label(),
                "energy": h.energy.as_ref().map_or(Value::Null, |(j, c)| json!({ "joules_micro": j, "class": c })),
                "authority": h.authority.label(),
                "synthesis": h.synthesis.map_or(Value::Null, |x| Value::from(x.label())),
            })).collect::<Vec<_>>(),
            "sources": self.sources.iter().map(|s| json!({
                "object": hex(&s.object),
                "signer": hex(&s.signer),
                "signer_id": s.signer_id,
            })).collect::<Vec<_>>(),
            "energy": self.energy.iter().map(|(c, j)| json!({ "class": c, "joules_micro": j })).collect::<Vec<_>>(),
            "unmetered": self.unmetered,
            "provenance": self.provenance.iter().map(|p| json!({
                "input": p.input,
                "kind": p.kind.label(),
                "digest": hex(&p.digest),
                "compartment": p.compartment.label(),
            })).collect::<Vec<_>>(),
        })
    }
}

/// The class labels in energy-measurement §2's order, strongest first.
const CLASS_ORDER: [&str; 5] = ["CalibratedInstrument", "OnChipCounter", "HwShunt", "ModelBased", "Estimator"];

/// Verify the derived object whose record is `records[subject]` (§7).
pub fn verify(held: &Held, subject: usize) -> Result<Facts, Refusal> {
    let mut v = Verifier::new(held);
    v.hop(subject, 1)?;
    let provenance = Record::read(&held.records[subject]).map(|r| r.provenance).unwrap_or_default();
    let subject_hash = v.hops.first().map(|h| h.record).unwrap_or_default();
    let mut totals: BTreeMap<&str, u64> = BTreeMap::new();
    let mut unmetered = 0;
    for h in &v.hops {
        match &h.energy {
            Some((j, c)) => {
                let c = CLASS_ORDER.iter().find(|x| **x == c.as_str()).copied().unwrap_or("Estimator");
                *totals.entry(c).or_default() += j;
            }
            None => unmetered += 1,
        }
    }
    let energy = CLASS_ORDER.iter().filter_map(|c| totals.get(c).map(|j| ((*c).to_owned(), *j))).collect();
    Ok(Facts { subject: subject_hash, hops: v.hops, sources: v.sources, energy, unmetered, provenance })
}

struct Verifier<'a> {
    held: &'a Held,
    by_hash: BTreeMap<[u8; 32], usize>,
    groups: Vec<([u8; 32], GroupReceipt)>,
    objects: Vec<ObjectReceipt>,
    /// Each record verified, with its origins.
    done: BTreeMap<usize, BTreeSet<String>>,
    hops: Vec<Hop>,
    sources: Vec<Source>,
}

/// An input's or output's receipts as §7.2 steps 3 and 4.1 find them.
enum Found<'a> {
    NotHeld,
    Invalid,
    Ok(&'a GroupReceipt, &'a ObjectReceipt),
}

impl<'a> Verifier<'a> {
    fn new(held: &'a Held) -> Verifier<'a> {
        let mut by_hash = BTreeMap::new();
        for (i, t) in held.records.iter().enumerate() {
            if let Some(h) = held_hash(t) {
                by_hash.entry(h).or_insert(i);
            }
        }
        let mut groups = Vec::new();
        let mut objects = Vec::new();
        for r in &held.receipts {
            match r.get("kind").and_then(Value::as_str) {
                Some("group") => {
                    if let Ok(g) = GroupReceipt::from_json(r) {
                        groups.push((g.link_hash(), g));
                    }
                }
                Some("object") => {
                    if let Ok(o) = ObjectReceipt::from_json(r) {
                        objects.push(o);
                    }
                }
                _ => {}
            }
        }
        Verifier { held, by_hash, groups, objects, done: BTreeMap::new(), hops: Vec::new(), sources: Vec::new() }
    }

    /// The group receipt of link hash `link` and its object receipt at `object_id`, which
    /// must verify and name `content`.
    fn receipts(&self, link: &[u8; 32], object_id: u64, content: &[u8; 32]) -> Found<'_> {
        let Some((_, g)) = self.groups.iter().find(|(l, _)| l == link) else { return Found::NotHeld };
        let cands: Vec<&ObjectReceipt> = self
            .objects
            .iter()
            .filter(|o| o.track_alias == g.track_alias && o.group_id == g.group_id && o.object_id == object_id)
            .collect();
        if cands.is_empty() {
            return Found::NotHeld;
        }
        if g.verify().is_err() {
            return Found::Invalid;
        }
        match cands.into_iter().find(|o| o.content_hash == *content && o.verify_inclusion(g).is_ok()) {
            Some(o) => Found::Ok(g, o),
            None => Found::Invalid,
        }
    }

    /// §7.2 for `records[r]`, reached at `depth` hops from the subject. Its origins (§8).
    fn hop(&mut self, r: usize, depth: usize) -> Result<BTreeSet<String>, Refusal> {
        if let Some(o) = self.done.get(&r) {
            return Ok(o.clone());
        }
        let held = self.held;
        let text = held.records.get(r).ok_or(Refusal { reason: Reason::HopNotHeld, at: Place::Record(r) })?;
        // Step 1, step 2.
        let rec = Record::read(text).map_err(|reason| Refusal { reason, at: Place::Record(r) })?;
        rec.verify_signature().map_err(|reason| Refusal { reason, at: Place::Record(r) })?;
        let me = self.hops.len();
        self.hops.push(Hop {
            record: rec.hash(),
            signer: rec.signer_pubkey,
            derivation: rec.derivation,
            check: rec.check,
            checked: Checked::Attested,
            energy: rec.energy.as_ref().map(|e| (e.joules_micro, e.class.label().to_owned())),
            authority: Authority::None,
            synthesis: rec.synthesis,
        });
        self.output(r, &rec)?;
        // Step 4.
        let mut origins = BTreeSet::new();
        let mut read: Vec<Option<Read>> = Vec::new();
        for (i, input) in rec.inputs.iter().enumerate() {
            let at = || Place::Input(r, i);
            let g = match self.receipts(&input.receipt, input.object_id, &input.content_hash) {
                Found::NotHeld => return refuse(Reason::InputReceiptNotHeld, at()),
                Found::Invalid => return refuse(Reason::InputReceiptInvalid, at()),
                Found::Ok(g, _) => g.clone(),
            };
            let mut below: Option<Record> = None;
            match input.derivation {
                None => {
                    if !held.sources.contains(&g.signer_pubkey) {
                        return refuse(Reason::ChainSkipsHop, at());
                    }
                    if let SignerIdForm::OriginBound { origin, .. } = classify(&g.signer_id) {
                        origins.insert(origin);
                    }
                    if !self.sources.iter().any(|s| s.object == input.content_hash) {
                        self.sources.push(Source {
                            object: input.content_hash,
                            signer: g.signer_pubkey,
                            signer_id: g.signer_id.clone(),
                        });
                    }
                }
                Some(h) => {
                    if depth + 1 > MAX_HOPS {
                        return refuse(Reason::ChainTooLong, at());
                    }
                    let Some(&j) = self.by_hash.get(&h) else { return refuse(Reason::HopNotHeld, at()) };
                    let below_rec =
                        Record::read(&held.records[j]).map_err(|reason| Refusal { reason, at: Place::Record(j) })?;
                    let o = &below_rec.output;
                    if o.content_hash != input.content_hash
                        || o.capability != input.capability
                        || o.receipt != input.receipt
                        || o.object_id != input.object_id
                    {
                        return refuse(Reason::HopMismatch, at());
                    }
                    origins.extend(self.hop(j, depth + 1)?);
                    below = Some(below_rec);
                }
            }
            let env = held.objects.get(&input.content_hash);
            read.push(match env {
                None => None,
                Some(env) => Some(read_input(&rec, i, env).map_err(|reason| Refusal { reason, at: at() })?),
            });
            // Step 4.5: the provenance it had, carried forward, and its label not lowered.
            let listed: BTreeSet<(ProvKind, [u8; 32])> =
                rec.provenance.iter().filter(|p| p.input == i as u64).map(|p| (p.kind, p.digest)).collect();
            let mut had: Vec<(ProvKind, [u8; 32])> = held
                .claims
                .iter()
                .filter_map(|t| claim_about(t))
                .filter(|(subject, _)| *subject == input.content_hash)
                .map(|(_, link)| (ProvKind::Claim, link))
                .collect();
            let stores = held.c2pa.get(&input.content_hash).into_iter().flatten().cloned();
            let embedded = env.map(|e| ops::embedded_c2pa(&input.capability, e)).unwrap_or_default();
            had.extend(stores.chain(embedded).map(|b| (ProvKind::C2pa, sha256(&b))));
            if had.iter().any(|x| !listed.contains(x)) {
                return refuse(Reason::ProvenanceStripped, at());
            }
            if let Some(b) = &below {
                let carried = b.provenance.iter().all(|p| {
                    rec.provenance
                        .iter()
                        .any(|q| q.input == i as u64 && q.kind == p.kind && q.digest == p.digest && q.compartment == p.compartment)
                });
                if !carried {
                    return refuse(Reason::ProvenanceStripped, at());
                }
                if input.synthesis < b.synthesis {
                    return refuse(Reason::SynthesisDropped, at());
                }
            }
        }
        // Step 5.
        let checked = self.recompute(r, &rec, &read)?;
        // Step 6.
        let authority = self.authority(r, &rec, &origins)?;
        // Step 7.
        if rec.check == Check::Attest && held.policy.refuse_attested {
            return refuse(Reason::AttestedRefused, Place::Policy(r));
        }
        self.hops[me].checked = checked;
        self.hops[me].authority = authority;
        self.done.insert(r, origins.clone());
        Ok(origins)
    }

    /// Step 3.
    fn output(&self, r: usize, rec: &Record) -> Result<(), Refusal> {
        let at = || Place::Output(r);
        let out = &rec.output;
        let (g, o) = match self.receipts(&out.receipt, out.object_id, &out.content_hash) {
            Found::NotHeld => return refuse(Reason::OutputReceiptNotHeld, at()),
            Found::Invalid => return refuse(Reason::OutputReceiptInvalid, at()),
            Found::Ok(g, o) => (g, o),
        };
        if g.signer_pubkey != rec.signer_pubkey {
            return refuse(Reason::OutputReceiptMismatch, at());
        }
        let figure = rec.energy.as_ref().map_or(0, |e| e.joules_micro);
        let label_agrees = match (&o.label, &rec.energy) {
            (None, _) => true,
            (Some(ObjectLabel::Unmetered), None) => true,
            (Some(ObjectLabel::Class(p)), Some(e)) => *p == e.class,
            _ => false,
        };
        if o.joules_micro != figure || !label_agrees {
            return refuse(Reason::EnergyDisagrees, at());
        }
        let Some(env) = self.held.objects.get(&out.content_hash) else { return Ok(()) };
        if content_hash(env) != out.content_hash {
            return refuse(Reason::OutputHashMismatch, at());
        }
        if rec.derivation.evaluates() {
            let (w, h) = (out.width.unwrap_or(0), out.height.unwrap_or(0));
            let Some((manifest, payload)) = ops::split_wai1(env) else {
                return refuse(Reason::OutputManifestMismatch, at());
            };
            if manifest != ops::output_manifest(rec, w, h).as_bytes() {
                return refuse(Reason::OutputManifestMismatch, at());
            }
            if out.capability == crate::codecs::CAP_IMAGE_PNG {
                let ok = crate::png_still::decode(payload).is_ok_and(|(pw, ph, rgb)| {
                    pw as u64 == w && ph as u64 == h && Some(super::sha256(&rgb)) == out.samples
                });
                if !ok {
                    return refuse(Reason::SamplesMismatch, at());
                }
            }
        } else if ops::wai1_capability(env).as_deref() != Some(out.capability.as_str()) {
            return refuse(Reason::OutputManifestMismatch, at());
        }
        Ok(())
    }

    /// Step 5.
    fn recompute(&self, r: usize, rec: &Record, read: &[Option<Read>]) -> Result<Checked, Refusal> {
        let at = || Place::Recompute(r);
        let all: Option<Vec<&Read>> = read.iter().map(Option::as_ref).collect();
        let Some(all) = all else {
            return match rec.check {
                Check::Attest => Ok(Checked::Attested),
                Check::Replicate if self.held.policy.require_recompute => refuse(Reason::InputsNotHeld, at()),
                Check::Replicate => Ok(Checked::NotRecomputed),
            };
        };
        if !rec.derivation.evaluates() {
            let Read::Exploded(env) = all[0] else { unreachable!("an explode reads an envelope") };
            if content_hash(env) != rec.output.content_hash {
                return refuse(Reason::ReplicateMismatch, at());
            }
            return Ok(Checked::Recomputed);
        }
        let samples: Vec<Samples> = all
            .iter()
            .map(|x| match x {
                Read::Samples(s) => s.clone(),
                Read::Exploded(_) => unreachable!("an evaluation reads samples"),
            })
            .collect();
        let exact = rec.derivation.exact(samples.len());
        if !exact {
            // An attested operation: nothing of it is pinned to recompute.
            return Ok(Checked::Attested);
        }
        let out = ops::evaluate(rec.derivation, &rec.params, &samples)
            .ok_or(Refusal { reason: Reason::ParamsInvalid, at: at() })?;
        let samples_agree = Some(out.digest()) == rec.output.samples
            && Some(out.w as u64) == rec.output.width
            && Some(out.h as u64) == rec.output.height;
        match rec.check {
            Check::Attest => {
                if !samples_agree {
                    return refuse(Reason::SamplesMismatch, at());
                }
                Ok(Checked::Attested)
            }
            Check::Replicate => {
                let manifest = ops::output_manifest(rec, out.w as u64, out.h as u64);
                let env = ops::wai1(manifest.as_bytes(), &crate::png_still::encode(&out.rgb, out.w, out.h));
                if !samples_agree || content_hash(&env) != rec.output.content_hash {
                    return refuse(Reason::ReplicateMismatch, at());
                }
                Ok(Checked::Recomputed)
            }
        }
    }

    /// Step 6 (§8).
    fn authority(&self, r: usize, rec: &Record, origins: &BTreeSet<String>) -> Result<Authority, Refusal> {
        let mut confirmed = true;
        for origin in origins {
            let at = || Place::Authority(r, Some(origin.clone()));
            let Some(text) = self.held.heads.get(origin) else {
                if self.held.policy.require_authority {
                    return refuse(Reason::EdgeNotAuthorised, at());
                }
                confirmed = false;
                continue;
            };
            let rev = Revision::from_json(text).map_err(|_| Refusal { reason: Reason::KeyLogNotHeld, at: at() })?;
            rev.check_held(origin).map_err(|_| Refusal { reason: Reason::KeyLogNotHeld, at: at() })?;
            let entry = rev
                .keys
                .iter()
                .find(|e| e.key == rec.signer_pubkey && e.may_sign.iter().any(|a| a == AUTHORITY))
                .ok_or(Refusal { reason: Reason::EdgeNotAuthorised, at: at() })?;
            if entry.status == EntryStatus::Revoked {
                return refuse(Reason::EdgeRevoked, at());
            }
        }
        if origins.is_empty() {
            if self.held.policy.require_authority {
                return refuse(Reason::EdgeNotAuthorised, Place::Authority(r, None));
            }
            return Ok(Authority::None);
        }
        Ok(if confirmed { Authority::Confirmed } else { Authority::Unconfirmed })
    }
}

/// What step 4.4 read of a held input: its samples, or for an explode the exploded
/// envelope.
enum Read {
    Samples(Samples),
    Exploded(Vec<u8>),
}

/// Step 4.4 for input `i` of `rec`, held as `env`.
fn read_input(rec: &Record, i: usize, env: &[u8]) -> Result<Read, Reason> {
    let input = &rec.inputs[i];
    if content_hash(env) != input.content_hash {
        return Err(Reason::InputHashMismatch);
    }
    if !rec.derivation.evaluates() {
        let Params::Explode { component } = rec.params else { unreachable!("an explode's parameters") };
        return ops::explode(env, component, &input.capability).map(Read::Exploded).ok_or(Reason::InputUnreadable);
    }
    let s = ops::read_samples(&input.capability, env, input.frame.unwrap_or(0)).ok_or(Reason::InputUnreadable)?;
    if Some(s.digest()) != input.samples {
        return Err(Reason::SamplesMismatch);
    }
    Ok(Read::Samples(s))
}
