//! Replay + live engines — extension §5.5–§5.6, §7.
//!
//! The normative core is convergence: any conformant sink holding the
//! same world container and the same op set MUST compute the identical
//! state hash at every tick, regardless of arrival order. This
//! reference implementation buys that property the simplest possible
//! way — ops live in a map keyed by their canonical order key, and a
//! late-arriving op triggers re-derivation from genesis. Production
//! sinks will keep snapshots and roll back incrementally; *how* a sink
//! converges is out of scope, *that* it converges is not, and this
//! engine is the oracle the conformance corpus is checked against.

use super::container::WorldContainer;
use super::graph::WorldState;
use super::op::{CausalOp, OpId};
use super::rng::Xoshiro256pp;
use super::sim::Sim;
use super::WorldError;
use std::collections::BTreeMap;

/// `(tick, state_hash)` at a contract checkpoint.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Checkpoint {
    pub tick: u64,
    pub hash: [u8; 32],
}

/// Canonical order key (extension §5.2).
type OrderKey = (u64, u64, [u8; 32]);

fn order_key(op: &CausalOp) -> OrderKey {
    (op.effect_tick, op.lamport, op.actor)
}

// ---------------------------------------------------------------------------
// Sealed replay
// ---------------------------------------------------------------------------

/// Replays sealed `wai.world.replay` containers.
pub struct Replayer<M: Sim> {
    pub sim: M,
}

impl<M: Sim> Replayer<M> {
    pub fn new(sim: M) -> Self {
        Replayer { sim }
    }

    /// Validate and replay a sealed log; returns the checkpoint list
    /// (every `checkpoint_every` ticks plus the final tick).
    ///
    /// Validation, all REQUIRED for sealed logs (extension §5.3, §7.2):
    /// signatures verify; `effect_tick >= 1`; ops listed in strictly
    /// ascending canonical order (which also forbids duplicates); every
    /// dep names an op present in the log; no op lands after the seal's
    /// final tick; and the replayed final state hash equals the seal's.
    /// A mismatch is a conformance failure — reported failed, never
    /// presented as approximately right.
    pub fn run_sealed(&self, container: &WorldContainer) -> Result<Vec<Checkpoint>, WorldError> {
        let seal = container.seal.ok_or(WorldError::BadSeal("replay requires a seal section"))?;

        let mut ids = std::collections::BTreeSet::new();
        let mut last_key: Option<OrderKey> = None;
        for op in &container.ops {
            op.verify()?;
            if op.effect_tick == 0 {
                return Err(WorldError::EffectTickZero);
            }
            if op.effect_tick > seal.final_tick {
                return Err(WorldError::BadSeal("op lands after the sealed final tick"));
            }
            let key = order_key(op);
            if let Some(prev) = last_key
                && key <= prev
            {
                return Err(WorldError::NotCanonicalOrder);
            }
            last_key = Some(key);
            if !ids.insert(op.id()) {
                return Err(WorldError::DuplicateOp(op.id()));
            }
        }
        for op in &container.ops {
            for dep in &op.deps {
                if !ids.contains(dep) {
                    return Err(WorldError::DepMissing(*dep));
                }
            }
        }

        let mut by_key: BTreeMap<OrderKey, &CausalOp> = BTreeMap::new();
        for op in &container.ops {
            by_key.insert(order_key(op), op);
        }
        let (checkpoints, _work) = simulate(&self.sim, container, &by_key, seal.final_tick)?;

        let final_hash = checkpoints
            .last()
            .map(|c| c.hash)
            .ok_or(WorldError::BadSeal("zero-tick replay"))?;
        if final_hash != seal.final_hash {
            return Err(WorldError::HashMismatch {
                tick: seal.final_tick,
                expected: seal.final_hash,
                got: final_hash,
            });
        }
        Ok(checkpoints)
    }
}

/// Shared tick loop: `state(t) = step(state(t-1) + ops_at(t))`, ops
/// within a tick applied in canonical order (which IS the map order).
fn simulate<M: Sim>(
    sim: &M,
    container: &WorldContainer,
    by_key: &BTreeMap<OrderKey, &CausalOp>,
    final_tick: u64,
) -> Result<(Vec<Checkpoint>, u64), WorldError> {
    let mut world: WorldState<M::Node> = WorldState::from_canonical_bytes(sim, &container.graph_bytes)?;
    let mut rng = Xoshiro256pp::from_seed(container.contract.rng_seed);
    let mut checkpoints = Vec::new();
    let cadence = container.contract.checkpoint_every;
    // Deterministic work: node-steps = Σ over ticks of the node count
    // integrated that tick. Exact and reproducible on any sink — the
    // honest, hardware-independent measure that goes in the receipt;
    // joules are derived from it downstream, never the other way round.
    let mut node_steps: u64 = 0;

    for tick in 1..=final_tick {
        for (_, op) in by_key.range((tick, 0, [0u8; 32])..=(tick, u64::MAX, [0xffu8; 32])) {
            if sim.resolve_kind(op.mutation_kind).is_none() {
                return Err(WorldError::UnknownMutationKind(op.mutation_kind));
            }
            sim.apply(&mut world, op)?;
        }
        sim.step(&mut world, &container.contract, tick, &mut rng);
        node_steps += world.nodes.len() as u64;
        if tick % cadence == 0 || tick == final_tick {
            checkpoints.push(Checkpoint { tick, hash: world.state_hash(sim) });
        }
    }
    Ok((checkpoints, node_steps))
}

/// Replay a container (sealed or not) and return its checkpoint list —
/// the `(tick, state_hash)` pairs §7 defines, which a session receipt
/// (§8) commits to. Unlike [`Replayer::run_sealed`] this does not
/// require or check a seal; it just simulates the ops the container
/// carries. The ops MUST already be in canonical order (they are, when
/// they come from a `LiveWorld::canonical_ops` or a sealed log).
pub fn collect_checkpoints<M: Sim>(
    sim: &M,
    container: &WorldContainer,
    final_tick: u64,
) -> Result<Vec<Checkpoint>, WorldError> {
    Ok(collect_checkpoints_and_work(sim, container, final_tick)?.0)
}

/// Like [`collect_checkpoints`] but also returns the deterministic
/// node-step work count — the honest, reproducible quantity a session
/// receipt (§8) carries in its signed payload.
pub fn collect_checkpoints_and_work<M: Sim>(
    sim: &M,
    container: &WorldContainer,
    final_tick: u64,
) -> Result<(Vec<Checkpoint>, u64), WorldError> {
    let mut by_key: BTreeMap<OrderKey, &CausalOp> = BTreeMap::new();
    for op in &container.ops {
        by_key.insert(order_key(op), op);
    }
    simulate(sim, container, &by_key, final_tick)
}

/// Replay an ordered container and return the state hash at each requested
/// tick — a sampler for consumers whose cadence is their own, not the
/// contract's `checkpoint_every`. A linear-film renderer uses this to read
/// the world state at each *frame* tick; the result is the same
/// replay-equivalent hash a full replay would compute at that tick, so the
/// film inherits the world's convergence guarantee for free.
///
/// `ticks` MUST be ascending (duplicates allowed — the same tick sampled
/// twice yields the same hash twice). Tick `0` is the genesis state (the
/// initial graph, before any step). The ops MUST already be in canonical
/// order (they are, from a sealed log or `LiveWorld::canonical_ops`).
pub fn collect_checkpoints_sampled<M: Sim>(
    sim: &M,
    container: &WorldContainer,
    ticks: &[u64],
) -> Result<Vec<Checkpoint>, WorldError> {
    let mut by_key: BTreeMap<OrderKey, &CausalOp> = BTreeMap::new();
    for op in &container.ops {
        by_key.insert(order_key(op), op);
    }
    let mut world: WorldState<M::Node> =
        WorldState::from_canonical_bytes(sim, &container.graph_bytes)?;
    let mut rng = Xoshiro256pp::from_seed(container.contract.rng_seed);
    let mut out = Vec::with_capacity(ticks.len());
    let mut want = ticks.iter().copied().peekable();
    // Genesis: sample the initial state before any tick steps.
    while want.peek() == Some(&0) {
        out.push(Checkpoint { tick: 0, hash: world.state_hash(sim) });
        want.next();
    }
    let max_tick = ticks.last().copied().unwrap_or(0);
    for tick in 1..=max_tick {
        for (_, op) in by_key.range((tick, 0, [0u8; 32])..=(tick, u64::MAX, [0xffu8; 32])) {
            if sim.resolve_kind(op.mutation_kind).is_none() {
                return Err(WorldError::UnknownMutationKind(op.mutation_kind));
            }
            sim.apply(&mut world, op)?;
        }
        sim.step(&mut world, &container.contract, tick, &mut rng);
        while want.peek() == Some(&tick) {
            out.push(Checkpoint { tick, hash: world.state_hash(sim) });
            want.next();
        }
    }
    Ok(out)
}

/// Replay a sealed container with the sim its contract names — the
/// dispatch a real sink performs. A sink that lacks the named sim
/// reports the world inert (the same "missing capability" semantics
/// as SPEC.md §4 step 6), so the caller can fall back per the
/// envelope's declared `fallback`.
pub fn run_sealed_dispatch(container: &WorldContainer) -> Result<Vec<Checkpoint>, WorldError> {
    use super::feed::{FeedState, CAP_SIM_FEED0};
    use super::sim::{Ballistic0, Kinematic0, CAP_SIM_BALLISTIC0, CAP_SIM_KINEMATIC0};
    use super::skeletal::{Skeletal0, CAP_SIM_SKELETAL0};
    use super::truss::{Truss0, CAP_SIM_TRUSS0};
    match container.contract.sim.as_str() {
        CAP_SIM_KINEMATIC0 => Replayer::new(Kinematic0).run_sealed(container),
        CAP_SIM_BALLISTIC0 => Replayer::new(Ballistic0).run_sealed(container),
        CAP_SIM_TRUSS0 => Replayer::new(Truss0).run_sealed(container),
        CAP_SIM_FEED0 => Replayer::new(FeedState).run_sealed(container),
        CAP_SIM_SKELETAL0 => Replayer::new(Skeletal0).run_sealed(container),
        other => Err(WorldError::UnknownSimCapability(other.into())),
    }
}

/// [`collect_checkpoints_sampled`] with the sim the contract names — the
/// dispatch a linear-film sink performs to read its world at frame ticks.
pub fn run_sampled_dispatch(
    container: &WorldContainer,
    ticks: &[u64],
) -> Result<Vec<Checkpoint>, WorldError> {
    use super::feed::{FeedState, CAP_SIM_FEED0};
    use super::sim::{Ballistic0, Kinematic0, CAP_SIM_BALLISTIC0, CAP_SIM_KINEMATIC0};
    use super::skeletal::{Skeletal0, CAP_SIM_SKELETAL0};
    use super::truss::{Truss0, CAP_SIM_TRUSS0};
    match container.contract.sim.as_str() {
        CAP_SIM_KINEMATIC0 => collect_checkpoints_sampled(&Kinematic0, container, ticks),
        CAP_SIM_BALLISTIC0 => collect_checkpoints_sampled(&Ballistic0, container, ticks),
        CAP_SIM_TRUSS0 => collect_checkpoints_sampled(&Truss0, container, ticks),
        CAP_SIM_FEED0 => collect_checkpoints_sampled(&FeedState, container, ticks),
        CAP_SIM_SKELETAL0 => collect_checkpoints_sampled(&Skeletal0, container, ticks),
        other => Err(WorldError::UnknownSimCapability(other.into())),
    }
}

// ---------------------------------------------------------------------------
// Live world
// ---------------------------------------------------------------------------

/// Outcome of [`LiveWorld::ingest`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Ingest {
    /// Admitted (deps satisfied). May have triggered re-derivation if
    /// it landed at or before the current tick.
    Applied,
    /// Held in the dep-gate buffer until its deps arrive.
    Buffered,
    /// Exact duplicate of an already-admitted op — idempotent no-op
    /// (transports redeliver; redelivery must not be an error).
    Duplicate,
}

/// A `wai.world.live` sink: ops arrive in any order over any transport;
/// the engine admits, dep-gates, and converges.
pub struct LiveWorld<M: Sim> {
    sim: M,
    container: WorldContainer,
    /// Admitted ops in canonical order (map order = application order).
    known: BTreeMap<OrderKey, CausalOp>,
    /// Op-id index for dep checks + equivocation detection.
    ids: BTreeMap<OpId, OrderKey>,
    /// Dep-gated buffer.
    pending: Vec<CausalOp>,
    /// Latest receipted checkpoint (extension §5.6): ops at or before
    /// this tick are immutable history and MUST be rejected.
    sealed_tick: u64,
    current_tick: u64,
    state: WorldState<M::Node>,
    /// Advances in lockstep with `state` — reset together on re-derive,
    /// so a sim that draws per tick stays exact across late-op rewinds.
    rng: Xoshiro256pp,
    dirty: bool,
}

impl<M: Sim> LiveWorld<M> {
    /// Open a live world from its container (seal section absent; any
    /// ops already in the container are ingested through the same gate
    /// arriving ops use).
    pub fn open(sim: M, container: WorldContainer) -> Result<Self, WorldError> {
        if container.seal.is_some() {
            return Err(WorldError::BadSeal("live world must not carry a seal"));
        }
        let state = WorldState::from_canonical_bytes(&sim, &container.graph_bytes)?;
        let rng = Xoshiro256pp::from_seed(container.contract.rng_seed);
        let mut lw = LiveWorld {
            sim,
            known: BTreeMap::new(),
            ids: BTreeMap::new(),
            pending: Vec::new(),
            sealed_tick: 0,
            current_tick: 0,
            state,
            rng,
            container,
            dirty: false,
        };
        let ops = std::mem::take(&mut lw.container.ops);
        for op in ops {
            lw.ingest(op)?;
        }
        Ok(lw)
    }

    /// Verify, gate, and admit one op (extension §5.6).
    pub fn ingest(&mut self, op: CausalOp) -> Result<Ingest, WorldError> {
        op.verify()?;
        if op.effect_tick == 0 {
            return Err(WorldError::EffectTickZero);
        }
        if op.effect_tick <= self.sealed_tick {
            return Err(WorldError::SealedHistory {
                effect_tick: op.effect_tick,
                sealed_tick: self.sealed_tick,
            });
        }
        if let Some(existing_key) = self.ids.get(&op.id()) {
            // Same id again: idempotent if byte-identical, equivocation
            // (same (actor, lamport), different content) otherwise.
            return if self.known.get(existing_key) == Some(&op) {
                Ok(Ingest::Duplicate)
            } else {
                Err(WorldError::DuplicateOp(op.id()))
            };
        }
        if !op.deps.iter().all(|d| self.ids.contains_key(d)) {
            self.pending.push(op);
            return Ok(Ingest::Buffered);
        }
        self.admit(op);
        self.drain_pending();
        Ok(Ingest::Applied)
    }

    fn admit(&mut self, op: CausalOp) {
        if op.effect_tick <= self.current_tick {
            self.dirty = true; // landed in the past → re-derive
        }
        self.ids.insert(op.id(), order_key(&op));
        self.known.insert(order_key(&op), op);
    }

    fn drain_pending(&mut self) {
        loop {
            let mut admitted_any = false;
            let mut still = Vec::with_capacity(self.pending.len());
            for op in std::mem::take(&mut self.pending) {
                if op.deps.iter().all(|d| self.ids.contains_key(d)) {
                    self.admit(op);
                    admitted_any = true;
                } else {
                    still.push(op);
                }
            }
            self.pending = still;
            if !admitted_any {
                break;
            }
        }
    }

    /// Advance (or re-derive) to `tick` and return the state hash.
    /// Rewinds (tick < current) and late ops both take the re-derive
    /// path: genesis forward, the convergence oracle.
    pub fn advance_to(&mut self, tick: u64) -> Result<[u8; 32], WorldError> {
        if self.dirty || tick < self.current_tick {
            self.state = WorldState::from_canonical_bytes(&self.sim, &self.container.graph_bytes)?;
            self.rng = Xoshiro256pp::from_seed(self.container.contract.rng_seed);
            self.current_tick = 0;
            self.dirty = false;
        }
        for t in (self.current_tick + 1)..=tick {
            let ops: Vec<CausalOp> = self
                .known
                .range((t, 0, [0u8; 32])..=(t, u64::MAX, [0xffu8; 32]))
                .map(|(_, op)| op.clone())
                .collect();
            for op in &ops {
                if self.sim.resolve_kind(op.mutation_kind).is_none() {
                    return Err(WorldError::UnknownMutationKind(op.mutation_kind));
                }
                self.sim.apply(&mut self.state, op)?;
            }
            self.sim.step(&mut self.state, &self.container.contract, t, &mut self.rng);
        }
        self.current_tick = self.current_tick.max(tick);
        Ok(self.state.state_hash(&self.sim))
    }

    /// Mark history up to `tick` as receipted (extension §8): a sealed
    /// checkpoint is immutable — later ops claiming earlier ticks are
    /// rejected at ingest.
    pub fn seal_checkpoint(&mut self, tick: u64) {
        self.sealed_tick = self.sealed_tick.max(tick);
    }

    /// Number of ops still dep-gated.
    pub fn pending_len(&self) -> usize {
        self.pending.len()
    }

    /// The world at the last `advance_to` tick — what a renderer draws.
    pub fn state(&self) -> &WorldState<M::Node> {
        &self.state
    }

    /// The tick `state()` reflects.
    pub fn current_tick(&self) -> u64 {
        self.current_tick
    }

    /// Admitted ops in canonical order — what a publisher seals into a
    /// `wai.world.replay`.
    pub fn canonical_ops(&self) -> Vec<CausalOp> {
        self.known.values().cloned().collect()
    }
}
