//! The reference edge's derive (`extensions/deriving-edge.md` §3, §6): read the inputs,
//! evaluate the derivation, write the output, publish it in a group of the edge's own track
//! under the edge's key, and sign its record. The `wai_edge` tool runs it over a bundle of
//! objects and receipts; the corpus runs it to make every hop it holds.
//!
//! It derives from what it is given and checks only that the inputs read as the derivation
//! needs: a caller verifies its inputs first (§3.3), as `wai_edge derive` does with
//! [`super::verify`].

use super::ops::{self, Samples};
use super::{
    check_of, ordered, sha256, Compartment, Derivation, Encoder, Energy, Input, Output, Params, ProvItem, ProvKind, Record,
    Synthesis,
};
use crate::codecs::{CAP_IMAGE_JPEG, CAP_IMAGE_PNG};
use crate::jwp_delivery::{GroupBuilder, GroupReceipt, ObjectLabel, ObjectReceipt};
use crate::jwp_hash::content_hash;
use ed25519_dalek::SigningKey;

/// An input as the edge holds it: its envelope, its delivery receipts, the record that
/// derived it where an edge did, and the provenance it came with.
pub struct InputObject<'a> {
    pub envelope: &'a [u8],
    /// The capability the derivation reads it as (§5.1).
    pub capability: &'a str,
    /// The frame it is read at; `None` for an explode.
    pub frame: Option<u64>,
    pub group: &'a GroupReceipt,
    pub object: &'a ObjectReceipt,
    pub record: Option<&'a Record>,
    /// The provenance items it came with apart from its envelope and its record (§6.1):
    /// the claims about it and the C2PA Manifest Stores bound to it, with the compartment the
    /// edge reads each to be in. Those embedded in its envelope, and those its record carries,
    /// are added by [`derive`].
    pub provenance: Vec<(ProvKind, [u8; 32], Compartment)>,
    /// The AI-origin label its provenance states (§6.2); its record's, where it has one, is
    /// taken when it is higher.
    pub synthesis: Option<Synthesis>,
}

/// Every provenance item of input `i` (§6.1): what it came with, the C2PA Manifest Stores
/// embedded in its envelope (system provenance), and every item its record carries.
pub fn provenance_of(i: usize, input: &InputObject<'_>) -> Vec<ProvItem> {
    let n = i as u64;
    let mut items: Vec<ProvItem> = input
        .provenance
        .iter()
        .map(|(kind, digest, compartment)| ProvItem { input: n, kind: *kind, digest: *digest, compartment: *compartment })
        .collect();
    for store in ops::embedded_c2pa(input.capability, input.envelope) {
        items.push(ProvItem { input: n, kind: ProvKind::C2pa, digest: sha256(&store), compartment: Compartment::System });
    }
    if let Some(r) = input.record {
        items.extend(r.provenance.iter().map(|p| ProvItem { input: n, ..p.clone() }));
    }
    items
}

/// Where and how the edge publishes the output (§6).
pub struct Publish<'a> {
    pub key: &'a SigningKey,
    pub signer_id: &'a str,
    pub track_alias: u64,
    pub group_id: u64,
    /// The group's start and end, and the record's `issued_at`.
    pub time: &'a str,
    /// The link hash of the edge's previous group on the track.
    pub parent: Option<[u8; 32]>,
    /// The object receipt's `origin` label.
    pub origin: &'a str,
}

/// What a derive makes: the output envelope, its group and object receipts, its record.
#[derive(Clone, Debug)]
pub struct Derived {
    pub envelope: Vec<u8>,
    pub group: GroupReceipt,
    pub object: ObjectReceipt,
    pub record: Record,
}

/// What the derivation computed, before it is published: the output envelope, and for an
/// evaluation its samples.
pub struct Written {
    pub envelope: Vec<u8>,
    pub samples: Option<Samples>,
}

/// The record of a derivation over `inputs`, its output fields still to be filled.
pub fn draft(
    d: Derivation,
    params: Params,
    encoder: Option<Encoder>,
    output_capability: &str,
    inputs: Vec<Input>,
    energy: Option<Energy>,
    issued_at: &str,
) -> Record {
    let check = check_of(d, inputs.len(), output_capability, encoder.as_ref());
    let inputs_synthesis = inputs.iter().filter_map(|i| i.synthesis).max();
    Record {
        derivation: d,
        params,
        encoder,
        check,
        inputs,
        output: Output {
            content_hash: [0; 32],
            capability: output_capability.to_owned(),
            receipt: [0; 32],
            object_id: 0,
            width: None,
            height: None,
            samples: None,
        },
        energy,
        provenance: Vec::new(),
        synthesis: inputs_synthesis,
        issued_at: issued_at.to_owned(),
        signer_pubkey: [0; 32],
        signer_id: String::new(),
        sig: [0; 64],
    }
}

/// Read each input as the derivation needs it, filling each input's `samples`: its samples
/// for an evaluation, its exploded component for an explode.
pub fn read_inputs(rec: &mut Record, envelopes: &[&[u8]]) -> Result<(Vec<Samples>, Option<Vec<u8>>), String> {
    if !rec.derivation.evaluates() {
        let Params::Explode { component } = rec.params else { return Err("an explode's parameters".into()) };
        let env = envelopes.first().ok_or("no input")?;
        let out = ops::explode(env, component, &rec.inputs[0].capability).ok_or("the component does not explode")?;
        return Ok((Vec::new(), Some(out)));
    }
    let mut samples = Vec::new();
    for (i, env) in envelopes.iter().enumerate() {
        let input = &mut rec.inputs[i];
        let s = ops::read_samples(&input.capability, env, input.frame.unwrap_or(0))
            .ok_or_else(|| format!("input {i} does not read as {} at frame {:?}", input.capability, input.frame))?;
        input.samples = Some(s.digest());
        samples.push(s);
    }
    Ok((samples, None))
}

/// Write `out`, the samples an evaluation produced, as `rec` states the output (§4.3),
/// filling the record's size and samples digest.
pub fn write(rec: &mut Record, out: &Samples) -> Result<Vec<u8>, String> {
    rec.output.width = Some(out.w as u64);
    rec.output.height = Some(out.h as u64);
    rec.output.samples = Some(out.digest());
    let payload = match rec.output.capability.as_str() {
        CAP_IMAGE_PNG => crate::png_still::encode(&out.rgb, out.w, out.h),
        CAP_IMAGE_JPEG => {
            let q = rec.encoder.as_ref().and_then(ops::jpeg_quality).ok_or("a JPEG encoder needs a quality")?;
            ops::write_jpeg(out, q).ok_or("the JPEG does not encode")?
        }
        c => return Err(format!("no writer for {c}")),
    };
    let manifest = ops::output_manifest(rec, out.w as u64, out.h as u64);
    Ok(ops::wai1(manifest.as_bytes(), &payload))
}

/// Evaluate `rec` over `envelopes` and write its output: the whole derivation, before it is
/// published.
pub fn compute(rec: &mut Record, envelopes: &[&[u8]]) -> Result<Written, String> {
    let (samples, exploded) = read_inputs(rec, envelopes)?;
    if let Some(env) = exploded {
        return Ok(Written { envelope: env, samples: None });
    }
    let out = ops::evaluate(rec.derivation, &rec.params, &samples).ok_or("the inputs do not satisfy the parameters")?;
    let envelope = write(rec, &out)?;
    Ok(Written { envelope, samples: Some(out) })
}

/// Publish `envelope` as the output of `rec` (§6): one object in a group of the edge's
/// track, its receipt labelled with the record's energy, and the record signed.
pub fn publish(mut rec: Record, envelope: Vec<u8>, p: &Publish) -> Derived {
    let hash = content_hash(&envelope);
    let (joules, label) = match &rec.energy {
        None => (0, ObjectLabel::Unmetered),
        Some(e) => (e.joules_micro, ObjectLabel::Class(e.class.clone())),
    };
    let mut b = GroupBuilder::default();
    b.add(hash, joules, p.origin, Some(label));
    let (group, mut objects) = b
        .finish(p.key, p.signer_id, p.track_alias, p.group_id, p.time, p.time, p.parent, true)
        .expect("one labelled object composes to a class");
    rec.output.content_hash = hash;
    rec.output.receipt = group.link_hash();
    rec.output.object_id = 0;
    rec.signer_id = p.signer_id.to_owned();
    let record = rec.sign(p.key);
    Derived { envelope, group, object: objects.remove(0), record }
}

/// The whole derive: §4's derivation over `inputs`, published under `p` (§6), carrying every
/// provenance item of every input forward (§6.1), and labelled no lower than any input or than
/// `own`, the label the derivation's own step warrants (§6.2).
#[allow(clippy::too_many_arguments)]
pub fn derive(
    d: Derivation,
    params: Params,
    encoder: Option<Encoder>,
    output_capability: &str,
    inputs: &[InputObject<'_>],
    energy: Option<Energy>,
    own: Option<Synthesis>,
    p: &Publish<'_>,
) -> Result<Derived, String> {
    let named: Vec<Input> = inputs
        .iter()
        .map(|i| Input {
            content_hash: i.object.content_hash,
            capability: i.capability.to_owned(),
            receipt: i.group.link_hash(),
            object_id: i.object.object_id,
            frame: i.frame,
            samples: None,
            derivation: i.record.map(Record::hash),
            synthesis: i.synthesis.max(i.record.and_then(|r| r.synthesis)),
        })
        .collect();
    let mut rec = draft(d, params, encoder, output_capability, named, energy, p.time);
    rec.synthesis = rec.synthesis.max(own);
    rec.provenance = ordered(inputs.iter().enumerate().flat_map(|(i, x)| provenance_of(i, x)).collect());
    let envelopes: Vec<&[u8]> = inputs.iter().map(|i| i.envelope).collect();
    let written = compute(&mut rec, &envelopes)?;
    let derived = publish(rec, written.envelope, p);
    // What it signed reads under §5.3: a derive that would sign a record a verifier
    // refuses at reading stops here.
    Record::read(&derived.record.to_json()).map_err(|r| format!("the record does not read: {r}"))?;
    Ok(derived)
}
