//! What a derivation computes (`extensions/deriving-edge.md` §4): reading an input as
//! samples (§4.1), the operations (§4.2), the output's pinned manifest and payload (§4.3),
//! and the explode (§4.4). An edge and a verifier run the same functions: what the edge
//! wrote is what the verifier recomputes.

use super::{sha256, Derivation, Encoder, Flip, Params, Record, DOMAIN_RENDER_SETTINGS};
use crate::claim_receipt::{hex, jcs};
use crate::codecs::{CAP_IMAGE_PNG, CAP_VIDEO_INT_MOTION};
use crate::container::WaiMulti;
use serde_json::Value;

/// An image as the canonical `RGB8` buffer (reconstruction-binding §2).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Samples {
    pub w: u32,
    pub h: u32,
    pub rgb: Vec<u8>,
}

impl Samples {
    /// The samples digest: SHA-256 over the buffer.
    pub fn digest(&self) -> [u8; 32] {
        sha256(&self.rgb)
    }

    fn at(&self, x: u32, y: u32) -> &[u8] {
        let i = 3 * (y as usize * self.w as usize + x as usize);
        &self.rgb[i..i + 3]
    }
}

/// A `WAI1` envelope's manifest and payload, when `env` is exactly one `WAI1` envelope:
/// the magic, the two length-prefixed parts, and nothing after the payload.
pub fn split_wai1(env: &[u8]) -> Option<(&[u8], &[u8])> {
    if env.len() < 12 || &env[..4] != crate::container::MAGIC {
        return None;
    }
    let ml = u32::from_le_bytes(env[4..8].try_into().ok()?) as usize;
    let p_at = 8usize.checked_add(ml)?;
    let pl = u32::from_le_bytes(env.get(p_at..p_at + 4)?.try_into().ok()?) as usize;
    let end = p_at.checked_add(4)?.checked_add(pl)?;
    if end != env.len() {
        return None;
    }
    Some((&env[8..p_at], &env[p_at + 4..end]))
}

/// A `WAI1` envelope of `manifest` and `payload` (SPEC §2).
pub fn wai1(manifest: &[u8], payload: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(12 + manifest.len() + payload.len());
    out.extend_from_slice(crate::container::MAGIC);
    out.extend_from_slice(&(manifest.len() as u32).to_le_bytes());
    out.extend_from_slice(manifest);
    out.extend_from_slice(&(payload.len() as u32).to_le_bytes());
    out.extend_from_slice(payload);
    out
}

/// The capability a `WAI1` envelope names, read by the reference reader.
pub fn wai1_capability(env: &[u8]) -> Option<String> {
    split_wai1(env)?;
    crate::container::Wai::from_bytes(env).ok().map(|w| w.capability().to_owned())
}

/// §4.1: input `env`, whose record says it is of `capability`, read at `frame` as samples.
/// `None` where it does not read so (`input_unreadable`).
pub fn read_samples(capability: &str, env: &[u8], frame: u64) -> Option<Samples> {
    let (_, payload) = split_wai1(env)?;
    if wai1_capability(env)? != capability {
        return None;
    }
    match capability {
        CAP_IMAGE_PNG => {
            if frame != 0 {
                return None;
            }
            let (w, h, rgb) = crate::png_still::decode(payload).ok()?;
            let kinds = crate::png_still::chunk_kinds(payload)?;
            if kinds.iter().any(|k| ![*b"IHDR", *b"IDAT", *b"IEND", super::C2PA_PNG_CHUNK].contains(k)) {
                return None;
            }
            Some(Samples { w, h, rgb })
        }
        CAP_VIDEO_INT_MOTION => {
            let v = crate::int_wire::parse_wiv1(payload, Some(0)).ok()?;
            let s = crate::int_video::VideoStream {
                n_frames: v.n_frames as usize,
                h: v.h as usize,
                w: v.w as usize,
                block: v.block as usize,
                q: v.q as i64,
                offset: v.offset,
                mv_precision: v.mv_precision,
                cdf: &v.cdf,
                mvs: &v.mvs,
                rans: v.rans,
            };
            let clip = crate::int_video::decode_typed(&s).ok()?;
            let rgb = clip.frames.into_iter().nth(usize::try_from(frame).ok()?)?;
            Some(Samples { w: v.w, h: v.h, rgb })
        }
        _ => None,
    }
}

/// The C2PA Manifest Stores embedded in an input (§6.1): for a `wai.image.png` envelope, the
/// data of each `caBX` chunk of its payload (C2PA 2.2 §A.3.2), in order. Read by the chunks'
/// lengths alone; a payload whose chunks do not walk embeds none here.
pub fn embedded_c2pa(capability: &str, env: &[u8]) -> Vec<Vec<u8>> {
    let Some((_, p)) = split_wai1(env) else { return Vec::new() };
    if capability != CAP_IMAGE_PNG || crate::png_still::chunk_kinds(p).is_none() {
        return Vec::new();
    }
    let mut out = Vec::new();
    let mut at = 8usize;
    while at + 12 <= p.len() {
        let len = u32::from_be_bytes(p[at..at + 4].try_into().expect("4")) as usize;
        if p[at + 4..at + 8] == super::C2PA_PNG_CHUNK {
            out.push(p[at + 8..at + 8 + len].to_vec());
        }
        at += 12 + len;
    }
    out
}

/// §4.4: component `component` of the `WAI2` envelope `env`, whose record says it is of
/// `capability`, exploded. `None` where it cannot be (`input_unreadable`).
pub fn explode(env: &[u8], component: u64, capability: &str) -> Option<Vec<u8>> {
    let w = WaiMulti::from_bytes(env).ok()?;
    let i = usize::try_from(component).ok()?;
    if w.manifest.renditions.get(i)?.capability != capability {
        return None;
    }
    WaiMulti::explode(env, i).ok()
}

/// §4.2: the operation `d` with `p` over `inputs`. `None` where the inputs cannot satisfy
/// the parameters (`params_invalid`). Scaling is the reference edge's nearest sample, which
/// no verifier holds an edge to.
pub fn evaluate(d: Derivation, p: &Params, inputs: &[Samples]) -> Option<Samples> {
    let out = match (d, p) {
        (Derivation::Identity, Params::None) => inputs.first()?.clone(),
        (Derivation::Crop | Derivation::RoiSelection, Params::Region { x, y, w, h }) => {
            let s = inputs.first()?;
            if x.checked_add(*w)? > s.w as u64 || y.checked_add(*h)? > s.h as u64 {
                return None;
            }
            let (x, y, w, h) = (*x as u32, *y as u32, *w as u32, *h as u32);
            let mut rgb = Vec::with_capacity(3 * w as usize * h as usize);
            for j in 0..h {
                for i in 0..w {
                    rgb.extend_from_slice(s.at(x + i, y + j));
                }
            }
            Samples { w, h, rgb }
        }
        (Derivation::Rotation, Params::Rotation { quarter_turns }) => {
            let s = inputs.first()?;
            let (sw, sh) = (s.w, s.h);
            let (w, h) = if *quarter_turns == 2 { (sw, sh) } else { (sh, sw) };
            let mut rgb = Vec::with_capacity(s.rgb.len());
            for j in 0..h {
                for i in 0..w {
                    let (x, y) = match quarter_turns {
                        1 => (j, sh - 1 - i),
                        2 => (sw - 1 - i, sh - 1 - j),
                        _ => (sw - 1 - j, i),
                    };
                    rgb.extend_from_slice(s.at(x, y));
                }
            }
            Samples { w, h, rgb }
        }
        (Derivation::Mirror, Params::Mirror { flip }) => {
            let s = inputs.first()?;
            let mut rgb = Vec::with_capacity(s.rgb.len());
            for j in 0..s.h {
                for i in 0..s.w {
                    let (x, y) = match flip {
                        Flip::LeftRight => (s.w - 1 - i, j),
                        Flip::TopBottom => (i, s.h - 1 - j),
                    };
                    rgb.extend_from_slice(s.at(x, y));
                }
            }
            Samples { w: s.w, h: s.h, rgb }
        }
        (Derivation::GridComposition, Params::Grid { columns, rows }) => {
            let first = inputs.first()?;
            let (cw, ch) = (first.w, first.h);
            if inputs.iter().any(|s| s.w != cw || s.h != ch) || inputs.len() as u64 != columns * rows {
                return None;
            }
            let (w64, h64) = (*columns * cw as u64, *rows * ch as u64);
            if !super::size_ok(w64, h64) {
                return None;
            }
            let (w, h) = (w64 as u32, h64 as u32);
            let mut rgb = Vec::with_capacity(3 * w as usize * h as usize);
            for j in 0..h {
                for i in 0..w {
                    let k = (j / ch) as usize * *columns as usize + (i / cw) as usize;
                    rgb.extend_from_slice(inputs[k].at(i % cw, j % ch));
                }
            }
            Samples { w, h, rgb }
        }
        (Derivation::OverlayComposition, Params::Overlay { width, height, fill, at }) => {
            if at.len() != inputs.len() {
                return None;
            }
            let (w, h) = (*width as u32, *height as u32);
            let mut rgb: Vec<u8> = fill.iter().copied().cycle().take(3 * w as usize * h as usize).collect();
            for (s, (x, y)) in inputs.iter().zip(at) {
                for j in 0..s.h as u64 {
                    for i in 0..s.w as u64 {
                        let (ox, oy) = (x + i, y + j);
                        if ox < *width && oy < *height {
                            let o = 3 * (oy as usize * w as usize + ox as usize);
                            rgb[o..o + 3].copy_from_slice(s.at(i as u32, j as u32));
                        }
                    }
                }
            }
            Samples { w, h, rgb }
        }
        (Derivation::Scaling, Params::Scaling { width, height }) => {
            let s = inputs.first()?;
            let (w, h) = (*width as u32, *height as u32);
            let mut rgb = Vec::with_capacity(3 * w as usize * h as usize);
            for j in 0..h as u64 {
                for i in 0..w as u64 {
                    let x = (i * s.w as u64 / *width) as u32;
                    let y = (j * s.h as u64 / *height) as u32;
                    rgb.extend_from_slice(s.at(x, y));
                }
            }
            Samples { w, h, rgb }
        }
        _ => return None,
    };
    Some(out)
}

/// `derived_from.settings` (§4.3): `BLAKE3("wai:render-settings\x01" ‖ JCS({"derivation",
/// "encoder", "frame", "params"}))`.
pub fn settings_digest(r: &Record) -> [u8; 32] {
    let encoder = r.encoder.as_ref().map_or(Value::Null, |e| serde_json::json!({ "id": e.id, "settings": e.settings }));
    let settings = serde_json::json!({
        "derivation": r.derivation.id(),
        "encoder": encoder,
        "frame": r.inputs.first().and_then(|i| i.frame),
        "params": r.params.to_value(),
    });
    let mut h = blake3::Hasher::new();
    h.update(DOMAIN_RENDER_SETTINGS);
    h.update(jcs(&settings).expect("settings are within the claim subset").as_bytes());
    *h.finalize().as_bytes()
}

/// `derived_from.renderer` (§4.3): the derivation and `@1` for a pinned writer, the
/// encoder's `id` otherwise.
pub fn renderer(r: &Record) -> String {
    match &r.encoder {
        None => format!("{}@1", r.derivation.id()),
        Some(e) => e.id.clone(),
    }
}

/// §4.3: the manifest of an evaluation's output of `width` × `height`, exactly.
pub fn output_manifest(r: &Record, width: u64, height: u64) -> String {
    let cap = &r.output.capability;
    let kind = cap.rsplit('.').next().unwrap_or(cap);
    let mut s = format!(
        r#"{{"wai":"1.0","media":"image","intent":"replicate","model_requirement":{{"capability":"{cap}"}},"conditioning":{{"kind":"{kind}"}}"#
    );
    if let [only] = r.inputs.as_slice() {
        s.push_str(&format!(
            r#","derived_from":{{"motif":"{}","motif_capability":"{}","renderer":"{}","settings":"{}"}}"#,
            hex(&only.content_hash),
            only.capability,
            renderer(r),
            hex(&settings_digest(r)),
        ));
    }
    s.push_str(&format!(
        r#","target":{{"frames":1,"h":{height},"w":{width}}}}}"#
    ));
    s
}

/// The quality the reference edge's JPEG encoder reads from its settings.
pub fn jpeg_quality(e: &Encoder) -> Option<u8> {
    e.settings.get("quality")?.as_u64().filter(|q| (1..=100).contains(q)).map(|q| q as u8)
}

/// The reference edge's JPEG writer: a baseline JPEG of `s` at `quality`. Not pinned by the
/// extension (§4.3): a record naming it is `attest`.
pub fn write_jpeg(s: &Samples, quality: u8) -> Option<Vec<u8>> {
    let (w, h) = (u16::try_from(s.w).ok()?, u16::try_from(s.h).ok()?);
    let mut out = Vec::new();
    jpeg_encoder::Encoder::new(&mut out, quality).encode(&s.rgb, w, h, jpeg_encoder::ColorType::Rgb).ok()?;
    Some(out)
}
