//! The WWLD payload — extension §4.
//!
//! One container serves all three capabilities; they differ only in
//! the op log's seal state: empty (`wai.world.scene`), sealed
//! (`wai.world.replay`), open (`wai.world.live`).
//!
//! ```text
//! "WWLD" | u16 n_sections | n × (u8 kind, u32 off, u32 len) | sections
//! ```
//! Offsets are relative to the start of the sections block. Unknown
//! section kinds are ignored (forward compatibility within MINOR
//! revisions of the extension).

use super::contract::SimContract;
use super::op::CausalOp;
use super::{DOMAIN_SEAL, WorldError};
use ed25519_dalek::{Signer, SigningKey, VerifyingKey};

pub const MAGIC_WWLD: &[u8; 4] = b"WWLD";

pub const SECTION_CONTRACT: u8 = 0x01;
pub const SECTION_GRAPH: u8 = 0x02;
pub const SECTION_ASSETS: u8 = 0x03;
pub const SECTION_OPLOG: u8 = 0x04;
pub const SECTION_SEAL: u8 = 0x05;

/// Asset table entry (extension §4.1): `asset_hash` is the BLAKE3-256
/// of a complete WAI envelope, domain `wai:asset\x01` — assets ARE WAI
/// objects, fetched and decoded by ordinary capability dispatch.
/// Presentation-only: assets never enter the state hash.
#[derive(Clone, Debug, PartialEq)]
pub struct AssetEntry {
    pub hash: [u8; 32],
    pub role: u8,
    pub name: String,
}

/// Seal section: present iff the container is a `wai.world.replay`.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Seal {
    pub final_tick: u64,
    pub final_hash: [u8; 32],
    pub signature: [u8; 64],
}

#[derive(Clone, Debug)]
pub struct WorldContainer {
    pub contract: SimContract,
    /// The exact contract bytes carried on the wire (the seal commits
    /// to these bytes, so they are preserved verbatim, not re-emitted).
    pub contract_bytes: Vec<u8>,
    /// Canonical tick-0 node serialization ([`super::WorldState::canonical_bytes`]).
    pub graph_bytes: Vec<u8>,
    pub assets: Vec<AssetEntry>,
    pub ops: Vec<CausalOp>,
    pub seal: Option<Seal>,
}

impl WorldContainer {
    pub fn to_bytes(&self) -> Vec<u8> {
        let oplog: Vec<u8> = self.ops.iter().flat_map(|o| o.to_bytes()).collect();
        let assets = encode_assets(&self.assets);

        let mut sections: Vec<(u8, Vec<u8>)> = vec![
            (SECTION_CONTRACT, self.contract_bytes.clone()),
            (SECTION_GRAPH, self.graph_bytes.clone()),
        ];
        if !self.assets.is_empty() {
            sections.push((SECTION_ASSETS, assets));
        }
        sections.push((SECTION_OPLOG, oplog));
        if let Some(seal) = &self.seal {
            let mut s = Vec::with_capacity(104);
            s.extend_from_slice(&seal.final_tick.to_le_bytes());
            s.extend_from_slice(&seal.final_hash);
            s.extend_from_slice(&seal.signature);
            sections.push((SECTION_SEAL, s));
        }

        let mut out = Vec::new();
        out.extend_from_slice(MAGIC_WWLD);
        out.extend_from_slice(&(sections.len() as u16).to_le_bytes());
        // Section table, then payloads packed 8-byte aligned.
        let table_len = sections.len() * 9;
        let mut off = 0u32;
        let mut offsets = Vec::with_capacity(sections.len());
        for (_, body) in &sections {
            offsets.push(off);
            let padded = body.len().div_ceil(8) * 8;
            off += padded as u32;
        }
        for ((kind, body), o) in sections.iter().zip(&offsets) {
            out.push(*kind);
            out.extend_from_slice(&o.to_le_bytes());
            out.extend_from_slice(&(body.len() as u32).to_le_bytes());
        }
        debug_assert_eq!(out.len(), 6 + table_len);
        for (_, body) in &sections {
            out.extend_from_slice(body);
            let pad = body.len().div_ceil(8) * 8 - body.len();
            out.extend_from_slice(&[0u8; 8][..pad]);
        }
        out
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<WorldContainer, WorldError> {
        if bytes.len() < 6 {
            return Err(WorldError::Truncated);
        }
        if &bytes[..4] != MAGIC_WWLD {
            return Err(WorldError::BadMagic);
        }
        let n = u16::from_le_bytes(bytes[4..6].try_into().unwrap()) as usize;
        let table_end = 6 + n * 9;
        if bytes.len() < table_end {
            return Err(WorldError::Truncated);
        }
        let payloads = &bytes[table_end..];

        let mut contract_bytes: Option<Vec<u8>> = None;
        let mut graph_bytes: Option<Vec<u8>> = None;
        let mut assets = Vec::new();
        let mut oplog: Option<Vec<u8>> = None;
        let mut seal: Option<Seal> = None;

        for i in 0..n {
            let e = &bytes[6 + i * 9..6 + (i + 1) * 9];
            let kind = e[0];
            let off = u32::from_le_bytes(e[1..5].try_into().unwrap()) as usize;
            let len = u32::from_le_bytes(e[5..9].try_into().unwrap()) as usize;
            let end = off.checked_add(len).ok_or(WorldError::Truncated)?;
            if end > payloads.len() {
                return Err(WorldError::Truncated);
            }
            let body = &payloads[off..end];
            match kind {
                SECTION_CONTRACT => contract_bytes = Some(body.to_vec()),
                SECTION_GRAPH => graph_bytes = Some(body.to_vec()),
                SECTION_ASSETS => assets = decode_assets(body)?,
                SECTION_OPLOG => oplog = Some(body.to_vec()),
                SECTION_SEAL => {
                    if body.len() != 8 + 32 + 64 {
                        return Err(WorldError::BadSeal("seal section must be 104 bytes"));
                    }
                    seal = Some(Seal {
                        final_tick: u64::from_le_bytes(body[..8].try_into().unwrap()),
                        final_hash: body[8..40].try_into().unwrap(),
                        signature: body[40..104].try_into().unwrap(),
                    });
                }
                _ => {} // unknown kinds ignored, by spec
            }
        }

        let contract_bytes = contract_bytes.ok_or(WorldError::MissingSection(SECTION_CONTRACT))?;
        let graph_bytes = graph_bytes.ok_or(WorldError::MissingSection(SECTION_GRAPH))?;
        let oplog = oplog.ok_or(WorldError::MissingSection(SECTION_OPLOG))?;
        let contract = SimContract::from_json(&contract_bytes)?;
        let ops = CausalOp::parse_log(&oplog)?;

        Ok(WorldContainer { contract, contract_bytes, graph_bytes, assets, ops, seal })
    }

    /// The seal's signed bytes (extension §4, seal section):
    /// `DOMAIN_SEAL || BLAKE3(contract) || BLAKE3(graph) || BLAKE3(oplog)
    ///  || final_tick_le || final_hash`.
    /// Hashing the three sections (rather than raw container bytes)
    /// makes the commitment independent of section-table order and
    /// padding — the seal binds content, not framing.
    pub fn seal_preimage(&self, final_tick: u64, final_hash: &[u8; 32]) -> Vec<u8> {
        let oplog: Vec<u8> = self.ops.iter().flat_map(|o| o.to_bytes()).collect();
        let mut out = DOMAIN_SEAL.to_vec();
        out.extend_from_slice(blake3::hash(&self.contract_bytes).as_bytes());
        out.extend_from_slice(blake3::hash(&self.graph_bytes).as_bytes());
        out.extend_from_slice(blake3::hash(&oplog).as_bytes());
        out.extend_from_slice(&final_tick.to_le_bytes());
        out.extend_from_slice(final_hash);
        out
    }

    /// Seal this container as a `wai.world.replay` under the
    /// publisher's key.
    pub fn seal(&mut self, publisher: &SigningKey, final_tick: u64, final_hash: [u8; 32]) {
        let sig = publisher.sign(&self.seal_preimage(final_tick, &final_hash));
        self.seal = Some(Seal { final_tick, final_hash, signature: sig.to_bytes() });
    }

    /// Verify the seal signature against the publisher's key. (Whether
    /// `final_hash` matches the replayed state is the replayer's job —
    /// this checks authenticity, replay checks correctness.)
    pub fn verify_seal(&self, publisher: &VerifyingKey) -> Result<Seal, WorldError> {
        let seal = self.seal.ok_or(WorldError::BadSeal("no seal section"))?;
        let msg = self.seal_preimage(seal.final_tick, &seal.final_hash);
        // The key's bytes as it was decoded from, so a non-canonical spelling the library
        // read is refused here like any other (SPEC §7 *Signatures*).
        crate::sig::verify_strict(publisher.as_bytes(), &msg, &seal.signature).map_err(|_| WorldError::BadSignature)?;
        Ok(seal)
    }
}

fn encode_assets(assets: &[AssetEntry]) -> Vec<u8> {
    let mut out = Vec::new();
    out.extend_from_slice(&(assets.len() as u32).to_le_bytes());
    for a in assets {
        out.extend_from_slice(&a.hash);
        out.push(a.role);
        let name = a.name.as_bytes();
        out.extend_from_slice(&(name.len() as u32).to_le_bytes());
        out.extend_from_slice(name);
    }
    out
}

fn decode_assets(bytes: &[u8]) -> Result<Vec<AssetEntry>, WorldError> {
    if bytes.len() < 4 {
        return Err(WorldError::Truncated);
    }
    let n = u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize;
    let mut pos = 4usize;
    // Reserve no more entries than the bytes left could hold: the count is the input's,
    // and a reservation sized by it alone let a few bytes ask for gigabytes.
    let mut out = Vec::with_capacity(n.min(bytes.len() / 37));
    for _ in 0..n {
        if pos + 37 > bytes.len() {
            return Err(WorldError::Truncated);
        }
        let hash: [u8; 32] = bytes[pos..pos + 32].try_into().unwrap();
        let role = bytes[pos + 32];
        let name_len = u32::from_le_bytes(bytes[pos + 33..pos + 37].try_into().unwrap()) as usize;
        pos += 37;
        if pos + name_len > bytes.len() {
            return Err(WorldError::Truncated);
        }
        let name = String::from_utf8(bytes[pos..pos + name_len].to_vec())
            .map_err(|_| WorldError::Truncated)?;
        pos += name_len;
        out.push(AssetEntry { hash, role, name });
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::world::sim::CAP_SIM_KINEMATIC0;
    use crate::world::{NUMERIC_FIXED64, RESOLVE_LWW};

    /// The seal's signature is verified strictly (SPEC §7 *Signatures*), under the
    /// publisher key's bytes as they were decoded: a key of small order, a non-canonical
    /// spelling of one, and an R of small order are each refused, though a lenient check
    /// accepts each.
    #[test]
    fn the_seal_is_verified_strictly() {
        let contract = SimContract {
            tick_rate: 60,
            numeric: NUMERIC_FIXED64.into(),
            rng_alg: "xoshiro256++".into(),
            rng_seed: [0x42; 32],
            sim: CAP_SIM_KINEMATIC0.into(),
            resolve_default: RESOLVE_LWW.into(),
            checkpoint_every: 100,
        };
        let publisher = SigningKey::from_bytes(&[9; 32]);
        let mut c = WorldContainer { contract_bytes: contract.to_canonical_json(), contract, graph_bytes: vec![], assets: vec![], ops: vec![], seal: None };
        c.seal(&publisher, 10, [7; 32]);
        crate::sig::assert_refuses_forgeries(
            &(c, publisher.verifying_key().to_bytes()),
            |g: &mut (WorldContainer, [u8; 32]), k, s| {
                g.1 = k;
                g.0.seal.as_mut().expect("sealed").signature = s;
            },
            |g| g.0.seal_preimage(10, &[7; 32]),
            |g| VerifyingKey::from_bytes(&g.1).is_ok_and(|k| g.0.verify_seal(&k).is_ok()),
        );
    }
}
