//! Exploding a `WAI2` component into its `WAI1` envelope, and packing `WAI1` envelopes
//! into a `WAI2` (SPEC §7.1, *Exploding a component*; *Refinement components*).
//!
//! A component and its exploded envelope are the same object for every rule that
//! names an object by content hash, so a refinement carried in a `WAI2` and the same
//! layer carried alone reach the same stages (`extensions/staged-delivery.md` §9).
//! Explode reads the `WAI2` manifest's members as the exact text they are in the
//! envelope (`serde_json::value::RawValue`), never re-serialised, so a component packed
//! with its manifest in `envelope_manifest` explodes back to the bytes it was packed
//! from. Identity is the entry's `id`, `sha256` and `prior` (SPEC §7.1, *Component
//! identity*): a refinement entry names its base and producers by `id`, and its
//! exploded envelope names them by the content hashes of their exploded envelopes.

use std::collections::BTreeMap;

use serde::Deserialize;
use serde_json::value::RawValue;
use serde_json::Value;

use crate::container::{is_component_id, manifest_bytes, Prev, PrevRun, RefForm, Refines, Wai, WaiMulti, MAGIC_V2};

fn bad(m: impl Into<String>) -> std::io::Error {
    std::io::Error::new(std::io::ErrorKind::InvalidData, m.into())
}

#[derive(Deserialize)]
struct RawV2<'a> {
    #[serde(borrow)]
    wai: &'a RawValue,
    #[serde(borrow)]
    media: &'a RawValue,
    #[serde(borrow)]
    intent: &'a RawValue,
    #[serde(borrow)]
    renditions: Vec<RawEntry<'a>>,
    #[serde(default, borrow)]
    target: Option<&'a RawValue>,
    #[serde(default, borrow)]
    security: Option<&'a RawValue>,
}

#[derive(Deserialize)]
struct RawEntry<'a> {
    #[serde(borrow)]
    capability: &'a RawValue,
    #[serde(borrow)]
    kind: &'a RawValue,
    #[serde(default, borrow)]
    prior: Option<&'a RawValue>,
    #[serde(default, borrow)]
    envelope_manifest: Option<&'a RawValue>,
}

/// A `WAI1` envelope with manifest text `manifest` and `payload`.
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
}

/// `refines` with each reference translated by `map`, written in the form
/// [`Refines::to_value`] gives it, without whitespace.
fn translate(r: &Refines, map: &dyn Fn(&str) -> Option<String>) -> std::io::Result<Refines> {
    let tr = |s: &str| map(s).ok_or_else(|| bad(format!("refines names '{s}', which no entry before it is")));
    Ok(Refines {
        base: tr(&r.base)?,
        prev: match &r.prev {
            Prev::One(p) => Prev::One(tr(p)?),
            Prev::Runs(runs) => Prev::Runs(runs.iter().map(|x| Ok(PrevRun { producer: tr(&x.producer)?, count: x.count })).collect::<std::io::Result<_>>()?),
        },
        layer: r.layer,
        span: r.span,
    })
}

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

impl WaiMulti {
    /// Every component of the `WAI2` envelope `envelope` as its exploded `WAI1`
    /// envelope, in table order (SPEC §7.1, *Exploding a component*): the payload is the
    /// entry's rendition bytes; the manifest is the entry's `envelope_manifest` exactly
    /// as written, or else
    /// `{"wai":W,"media":M,"intent":I,"model_requirement":{"capability":C,"fallback":null},`
    /// `"conditioning":{"kind":K[,"prior":P]}[,"security":S][,"refines":R][,"target":T]}`,
    /// where `W`, `M`, `I`, `S` and `T` are the `WAI2` manifest's members as written, `C`,
    /// `K` and `P` the entry's, and `R` the entry's `refines` with each component id
    /// replaced by the content hash of that entry's exploded envelope. An entry a
    /// refinement names must come before it.
    pub fn explode_all(envelope: &[u8]) -> std::io::Result<Vec<Vec<u8>>> {
        Self::explode_each(envelope)?.into_iter().collect()
    }

    /// [`Self::explode_all`], entry by entry: an entry whose `refines` cannot be
    /// translated (it is malformed, or names an entry that is not before it, or one
    /// that could not be exploded) has its own error, and the others their envelopes.
    /// Entry `i`'s envelope is defined by entry `i` and the entries it names, so one
    /// entry that cannot be exploded leaves every entry that does not name it intact.
    pub fn explode_each(envelope: &[u8]) -> std::io::Result<Vec<std::io::Result<Vec<u8>>>> {
        let w = WaiMulti::from_bytes(envelope)?;
        let mtext = manifest_bytes(envelope)?;
        let raw: RawV2 = serde_json::from_slice(mtext).map_err(|e| bad(format!("WAI2 manifest: {e}")))?;
        if raw.renditions.len() != w.manifest.renditions.len() {
            return Err(bad("the manifest's entries do not read alike"));
        }
        let mut ids: BTreeMap<String, String> = BTreeMap::new();
        let mut out = Vec::with_capacity(raw.renditions.len());
        for (i, (e, meta)) in raw.renditions.iter().zip(&w.manifest.renditions).enumerate() {
            let payload = &w.renditions[i];
            let manifest = (|| -> std::io::Result<Vec<u8>> { Ok(match e.envelope_manifest {
                Some(m) => m.get().as_bytes().to_vec(),
                None => {
                    let mut s = String::new();
                    s.push_str("{\"wai\":");
                    s.push_str(raw.wai.get());
                    s.push_str(",\"media\":");
                    s.push_str(raw.media.get());
                    s.push_str(",\"intent\":");
                    s.push_str(raw.intent.get());
                    s.push_str(",\"model_requirement\":{\"capability\":");
                    s.push_str(e.capability.get());
                    s.push_str(",\"fallback\":null},\"conditioning\":{\"kind\":");
                    s.push_str(e.kind.get());
                    if let Some(p) = e.prior {
                        s.push_str(",\"prior\":");
                        s.push_str(p.get());
                    }
                    s.push('}');
                    if let Some(sec) = raw.security {
                        s.push_str(",\"security\":");
                        s.push_str(sec.get());
                    }
                    if let Some(v) = meta.refines.as_ref() {
                        let r = Refines::from_value(v, RefForm::ComponentId).map_err(|e| bad(format!("entry {i}: {e}")))?;
                        let t = translate(&r, &|id| ids.get(id).cloned())?;
                        s.push_str(",\"refines\":");
                        s.push_str(&serde_json::to_string(&t.to_value()).map_err(|e| bad(e.to_string()))?);
                    }
                    if let Some(t) = raw.target {
                        s.push_str(",\"target\":");
                        s.push_str(t.get());
                    }
                    s.push('}');
                    s.into_bytes()
                }
            }) })();
            let env = manifest.map(|m| wai1(&m, payload));
            if let (Some(id), Ok(env)) = (meta.id.as_deref(), env.as_ref()) {
                ids.insert(id.to_string(), hex(&crate::jwp_hash::content_hash(env)));
            }
            out.push(env);
        }
        Ok(out)
    }

    /// Component `index` of the `WAI2` envelope `envelope` as its exploded `WAI1`
    /// envelope ([`Self::explode_each`]): its own result, whatever the other entries'.
    pub fn explode(envelope: &[u8], index: usize) -> std::io::Result<Vec<u8>> {
        Self::explode_each(envelope)?.into_iter().nth(index).ok_or_else(|| bad(format!("no component {index}")))?
    }

    /// Pack `WAI1` envelopes into one `WAI2`, in order: each entry carries the
    /// envelope's capability, kind and pin, the id `c<i>`, its payload's `sha256`, and
    /// its manifest verbatim as `envelope_manifest`; an entry whose capability is
    /// refinement-only takes `role: "refinement"` and its `refines` with each content
    /// hash replaced by the id of the entry whose envelope has it. `wai`, `media`,
    /// `intent` and `target` come from the first envelope. So every envelope explodes
    /// back to its own bytes ([`Self::explode_all`]). Refused: an encrypted envelope;
    /// a refinement whose base or producers are not among the envelopes before it; an
    /// envelope with bytes around its manifest's JSON value or after its payload, which
    /// no explode writes back (SPEC §7.1); and an output that does not parse, explode
    /// back to every source byte for byte, and pass
    /// [`Self::check_envelope_manifests`], which pack checks before it returns (a
    /// manifest nested near the JSON depth limit as a `WAI1` passes it as a member of
    /// a `WAI2` manifest).
    pub fn pack(w1s: &[&[u8]]) -> std::io::Result<Vec<u8>> {
        let out = Self::pack_unchecked(w1s)?;
        WaiMulti::from_bytes(&out).map_err(|e| bad(format!("the packed WAI2 does not parse: {e}")))?;
        let back = Self::explode_all(&out).map_err(|e| bad(format!("the packed WAI2 does not explode: {e}")))?;
        if back.len() != w1s.len() || back.iter().zip(w1s).any(|(a, b)| a.as_slice() != *b) {
            return Err(bad("the packed WAI2 does not explode back to its sources"));
        }
        Self::check_envelope_manifests(&out).map_err(|(i, c)| bad(format!("the packed WAI2's entry {i}: {c}")))?;
        Ok(out)
    }

    fn pack_unchecked(w1s: &[&[u8]]) -> std::io::Result<Vec<u8>> {
        if w1s.is_empty() {
            return Err(bad("nothing to pack"));
        }
        let mut entries = Vec::new();
        let mut payloads: Vec<Vec<u8>> = Vec::new();
        let mut by_hash: BTreeMap<String, String> = BTreeMap::new();
        let mut first: Option<(String, String, String)> = None;
        for (i, b) in w1s.iter().enumerate() {
            let w = Wai::from_bytes(b)?;
            if w.manifest.security.as_ref().is_some_and(|s| !s.is_plaintext()) {
                return Err(bad(format!("envelope {i} is encrypted")));
            }
            let mtext = std::str::from_utf8(manifest_bytes(b)?).map_err(|e| bad(e.to_string()))?;
            // The manifest is copied as the value of `envelope_manifest`, where bytes
            // around a JSON value are not part of it: such a manifest cannot come back.
            let value: &RawValue = serde_json::from_str(mtext).map_err(|e| bad(e.to_string()))?;
            if value.get().len() != mtext.len() {
                return Err(bad(format!("envelope {i}'s manifest has bytes around its JSON value")));
            }
            if b.len() != 12 + mtext.len() + w.payload.len() {
                return Err(bad(format!("envelope {i} has bytes after its payload")));
            }
            let id = format!("c{i}");
            debug_assert!(is_component_id(&id));
            let mut e = serde_json::Map::new();
            e.insert("capability".into(), Value::from(w.capability()));
            e.insert("kind".into(), Value::from(w.kind()));
            e.insert("id".into(), Value::from(id.clone()));
            e.insert("sha256".into(), Value::from(crate::pin::sha256_hex(&w.payload)));
            if let Some(p) = &w.manifest.conditioning.prior {
                e.insert("prior".into(), serde_json::to_value(p).map_err(|x| bad(x.to_string()))?);
            }
            if crate::codecs::is_refinement_only(w.capability()) {
                e.insert("role".into(), Value::from("refinement"));
                let r = w.refines().ok_or_else(|| bad(format!("envelope {i} is a refinement with no refines")))?.map_err(|x| bad(x.to_string()))?;
                let t = translate(&r, &|h| by_hash.get(h).cloned())?;
                e.insert("refines".into(), t.to_value());
            }
            let mut entry = serde_json::to_string(&Value::Object(e)).map_err(|x| bad(x.to_string()))?;
            entry.pop();
            entry.push_str(",\"envelope_manifest\":");
            entry.push_str(mtext);
            entry.push('}');
            entries.push(entry);
            if first.is_none() {
                let tv = |v: &Value| serde_json::to_string(v).unwrap_or_else(|_| "null".into());
                first = Some((tv(&Value::from(w.media())), tv(&Value::from(w.intent())), tv(&w.manifest.target)));
            }
            by_hash.insert(hex(&crate::jwp_hash::content_hash(b)), id);
            payloads.push(w.payload);
        }
        let (media, intent, target) = first.expect("one envelope");
        let manifest = format!("{{\"wai\":\"1.1\",\"media\":{media},\"intent\":{intent},\"renditions\":[{}],\"target\":{target}}}", entries.join(","));
        if payloads.len() > u16::MAX as usize {
            return Err(bad("too many components"));
        }
        let mut out = Vec::new();
        out.extend_from_slice(MAGIC_V2);
        out.extend_from_slice(&(manifest.len() as u32).to_le_bytes());
        out.extend_from_slice(manifest.as_bytes());
        out.extend_from_slice(&(payloads.len() as u16).to_le_bytes());
        let mut off = 0u32;
        for p in &payloads {
            out.extend_from_slice(&off.to_le_bytes());
            out.extend_from_slice(&(p.len() as u32).to_le_bytes());
            off = off.checked_add(p.len() as u32).ok_or_else(|| bad("payloads past u32"))?;
        }
        for p in &payloads {
            out.extend_from_slice(p);
        }
        Ok(out)
    }

    /// Check every entry's `envelope_manifest` against the entry (SPEC §7.1): it is a
    /// `WAI1` manifest whose `model_requirement.capability`, `conditioning.kind` and
    /// `conditioning.prior` are the entry's, whose `refines` is the entry's with each
    /// id replaced by the content hash of that entry's exploded envelope (or both are
    /// absent), and, for a refinement-only capability, whose `fallback` is `null`. The
    /// first entry that fails, with `component-manifest-mismatch`.
    pub fn check_envelope_manifests(envelope: &[u8]) -> Result<(), (usize, &'static str)> {
        const CODE: &str = "component-manifest-mismatch";
        let w = WaiMulti::from_bytes(envelope).map_err(|_| (0, CODE))?;
        let exploded = Self::explode_each(envelope).map_err(|_| (0, CODE))?;
        let mut ids: BTreeMap<String, String> = BTreeMap::new();
        for (i, (meta, env)) in w.manifest.renditions.iter().zip(&exploded).enumerate() {
            if let Some(m) = &meta.envelope_manifest {
                let mm: crate::container::Manifest = serde_json::from_value(m.clone()).map_err(|_| (i, CODE))?;
                let as_json = |p: Option<&crate::container::Prior>| p.map(|p| serde_json::to_value(p).ok());
                let prior_eq = as_json(mm.conditioning.prior.as_ref()) == as_json(meta.prior.as_ref());
                if mm.model_requirement.capability != meta.capability || mm.conditioning.kind != meta.kind || !prior_eq {
                    return Err((i, CODE));
                }
                // A refinement-only capability's manifest names no fallback
                // (staged-delivery §3).
                if crate::codecs::is_refinement_only(&meta.capability) && mm.model_requirement.fallback.is_some() {
                    return Err((i, CODE));
                }
                let entry_refines = match &meta.refines {
                    None => None,
                    Some(v) => {
                        let r = Refines::from_value(v, RefForm::ComponentId).map_err(|_| (i, CODE))?;
                        Some(translate(&r, &|id| ids.get(id).cloned()).map_err(|_| (i, CODE))?)
                    }
                };
                let wai1_refines = match &mm.refines {
                    None => None,
                    Some(v) => Some(Refines::from_value(v, RefForm::ContentHash).map_err(|_| (i, CODE))?),
                };
                if entry_refines != wai1_refines {
                    return Err((i, CODE));
                }
            }
            if let (Some(id), Ok(env)) = (meta.id.as_deref(), env.as_ref()) {
                ids.insert(id.to_string(), hex(&crate::jwp_hash::content_hash(env)));
            }
        }
        Ok(())
    }
}
