//! World state + state hashing — extension §7.1.
//!
//! State is a map from 128-bit node id to sim-owned node state. The
//! canonical serialization walks nodes in node-id order; per-node
//! field layout belongs to the sim capability's registration (the sim
//! supplies the serializer). Assets and presentation data never enter
//! this hash — that is what keeps conformance state-level while
//! presentation legitimately differs between sinks.

use super::sim::Sim;
use super::{DOMAIN_STATE, WorldError};
use std::collections::BTreeMap;

/// The world at one tick: nodes in id order (BTreeMap = canonical
/// iteration for free).
#[derive(Clone, Debug, Default, PartialEq)]
pub struct WorldState<S> {
    pub nodes: BTreeMap<[u8; 16], S>,
}

impl<S> WorldState<S> {
    pub fn new() -> Self {
        WorldState { nodes: BTreeMap::new() }
    }

    /// Canonical state bytes:
    /// `n_nodes u32 | n × (node_id 16 | state_len u32 | state bytes)`,
    /// nodes ascending by id. This is also the WWLD graph section's
    /// format (the tick-0 graph IS canonical state bytes at tick 0).
    pub fn canonical_bytes<M: Sim<Node = S>>(&self, sim: &M) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend_from_slice(&(self.nodes.len() as u32).to_le_bytes());
        let mut buf = Vec::new();
        for (id, node) in &self.nodes {
            out.extend_from_slice(id);
            buf.clear();
            sim.serialize_node(node, &mut buf);
            out.extend_from_slice(&(buf.len() as u32).to_le_bytes());
            out.extend_from_slice(&buf);
        }
        out
    }

    /// `BLAKE3(domain || canonical_state_bytes)` — the quantity
    /// replay-equivalence is defined over.
    pub fn state_hash<M: Sim<Node = S>>(&self, sim: &M) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(DOMAIN_STATE);
        h.update(&self.canonical_bytes(sim));
        *h.finalize().as_bytes()
    }

    /// Parse canonical state bytes back into a world (graph section
    /// load). Node order in the bytes MUST be ascending; rejecting
    /// disorder keeps the graph section unambiguous.
    pub fn from_canonical_bytes<M: Sim<Node = S>>(sim: &M, bytes: &[u8]) -> Result<Self, 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;
        let mut nodes = BTreeMap::new();
        let mut last_id: Option<[u8; 16]> = None;
        for _ in 0..n {
            if pos + 20 > bytes.len() {
                return Err(WorldError::Truncated);
            }
            let id: [u8; 16] = bytes[pos..pos + 16].try_into().unwrap();
            if let Some(prev) = last_id
                && id <= prev
            {
                return Err(WorldError::NotCanonicalOrder);
            }
            last_id = Some(id);
            pos += 16;
            let len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
            pos += 4;
            if pos + len > bytes.len() {
                return Err(WorldError::Truncated);
            }
            let node = sim.parse_node(&bytes[pos..pos + len])?;
            pos += len;
            nodes.insert(id, node);
        }
        if pos != bytes.len() {
            return Err(WorldError::Truncated);
        }
        Ok(WorldState { nodes })
    }
}
