//! Branching session state: an index arena of fixed-size entry records over
//! one contiguous payload arena, plus restoration from durable frames
//! (`SMH-SPEC-SPEC0001`, Sessions).
//!
//! Entry content lives as canonical CBOR bytes, at rest and in memory. A
//! record carries what traversal needs ([`EntryKind`]); anything else is a
//! decode into the consuming scope.

use crate::frame::Frame;
use serde::{Deserialize, Serialize};
use smith::error::{Result, SmithError};
use smith::id::{BranchId, EntryId, MessageId, SessionId, ToolCallId};
use smith::message::{ContentBlock, Message, Role};
use std::collections::HashMap;

/// Meta entry kind that marks a compaction boundary.
pub const COMPACTION_KIND: &str = "compaction";

/// Metadata kind for one redacted provider HTTP effect record.
pub const PROVIDER_HTTP_KIND: &str = "provider-http";

/// Deterministic JSON: recursively sort object keys so equivalent values
/// encode to identical CBOR regardless of insertion order.
fn canonical_json(value: &serde_json::Value) -> serde_json::Value {
    match value {
        serde_json::Value::Object(map) => {
            let mut keys: Vec<&String> = map.keys().collect();
            keys.sort();
            let mut out = serde_json::Map::new();
            for key in keys {
                out.insert(key.clone(), canonical_json(&map[key]));
            }
            serde_json::Value::Object(out)
        }
        serde_json::Value::Array(items) => {
            serde_json::Value::Array(items.iter().map(canonical_json).collect())
        }
        other => other.clone(),
    }
}

/// Whether `value` already has the key order [`canonical_json`] produces.
fn is_canonical(value: &serde_json::Value) -> bool {
    match value {
        serde_json::Value::Object(map) => {
            map.keys().zip(map.keys().skip(1)).all(|(a, b)| a < b) && map.values().all(is_canonical)
        }
        serde_json::Value::Array(items) => items.iter().all(is_canonical),
        _ => true,
    }
}

/// Durable session identity, persisted as the first frame of a session file
/// so reload reproduces the same IDs instead of minting new ones.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct SessionHeader {
    /// Session identifier.
    pub session_id: SessionId,
    /// Identifier of the root branch.
    pub branch_id: BranchId,
}

/// Byte range of one entry's content in its session's payload arena.
///
/// `u32` offsets cap one session arena at 4 GiB.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Span {
    /// First content byte.
    pub offset: u32,
    /// Content length in bytes.
    pub len: u32,
}

impl Span {
    /// The span of `len` bytes at `offset`, or an error past the arena cap.
    fn at(offset: usize, len: usize) -> Result<Self> {
        let full = || SmithError::Session {
            code: "SESSION_ARENA_FULL".to_string(),
            message: format!("entry content at {offset}+{len} exceeds the 4 GiB session arena"),
        };
        let end = u32::try_from(offset.saturating_add(len)).map_err(|_| full())?;
        let offset = u32::try_from(offset).map_err(|_| full())?;
        Ok(Self {
            offset,
            len: end - offset,
        })
    }

    const fn end(self) -> usize {
        self.offset as usize + self.len as usize
    }

    const fn range(self) -> std::ops::Range<usize> {
        self.offset as usize..self.end()
    }
}

/// What an entry recorded, reduced to the identifiers traversal pairs on.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum EntryKind {
    /// Chat message.
    Message {
        /// Message identifier.
        id: MessageId,
        /// Who produced the message.
        role: Role,
    },
    /// Tool call.
    ToolCall {
        /// Identifier pairing the call with its result.
        call_id: ToolCallId,
    },
    /// Tool result.
    ToolResult {
        /// Identifier of the call this result answers.
        call_id: ToolCallId,
        /// Whether the effect succeeded.
        ok: bool,
    },
    /// Session metadata.
    Meta {
        /// Whether the entry marks a compaction boundary.
        compaction: bool,
    },
}

/// A fixed-size session entry record; its content lives in the session
/// payload arena.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Entry {
    /// Entry ID.
    pub id: EntryId,
    /// Parent entry ID (None for root).
    pub parent: Option<EntryId>,
    /// Timestamp (Unix ms).
    pub timestamp_ms: i64,
    /// What the entry recorded.
    pub kind: EntryKind,
    /// Canonical CBOR of the [`EntryContent`], in the session arena.
    pub content: Span,
}

impl Entry {
    /// Whether this entry marks a compaction boundary.
    #[must_use]
    pub const fn is_compaction(&self) -> bool {
        matches!(self.kind, EntryKind::Meta { compaction: true })
    }
}

/// Wire shape of one entry frame: the record plus its content as an opaque
/// canonical CBOR byte string.
#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub struct EntryFrame {
    /// Entry ID.
    pub id: EntryId,
    /// Parent entry ID (None for root).
    pub parent: Option<EntryId>,
    /// Timestamp (Unix ms).
    pub timestamp_ms: i64,
    /// What the entry recorded.
    pub kind: EntryKind,
    /// Canonical CBOR of the [`EntryContent`].
    #[serde(serialize_with = "cbor_bytes", deserialize_with = "byte_buf")]
    pub content: Vec<u8>,
}

impl EntryFrame {
    /// A frame for `content` with a fresh identity, outside any session.
    ///
    /// # Errors
    ///
    /// Returns [`SmithError::Session`] when the content cannot be encoded.
    pub fn new(parent: Option<EntryId>, content: &EntryContent) -> Result<Self> {
        let mut bytes = Vec::new();
        content.encode_into(&mut bytes)?;
        Ok(Self {
            id: EntryId::new(),
            parent,
            timestamp_ms: now_ms(),
            kind: content.kind(),
            content: bytes,
        })
    }
}

/// Borrowed [`EntryFrame`] serialization of a session record and its arena
/// bytes, so arena content is framed without being copied out first.
#[derive(Serialize)]
pub(crate) struct EntryWire<'a> {
    id: EntryId,
    parent: Option<EntryId>,
    timestamp_ms: i64,
    kind: EntryKind,
    #[serde(serialize_with = "cbor_bytes")]
    content: &'a [u8],
}

impl<'a> EntryWire<'a> {
    pub(crate) const fn of(entry: &Entry, content: &'a [u8]) -> Self {
        Self {
            id: entry.id,
            parent: entry.parent,
            timestamp_ms: entry.timestamp_ms,
            kind: entry.kind,
            content,
        }
    }
}

/// Serialize bytes as one CBOR byte string rather than an integer array.
fn cbor_bytes<S: serde::Serializer>(
    bytes: &impl AsRef<[u8]>,
    serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
    serializer.serialize_bytes(bytes.as_ref())
}

fn byte_buf<'de, D: serde::Deserializer<'de>>(
    deserializer: D,
) -> std::result::Result<Vec<u8>, D::Error> {
    struct Visitor;
    impl serde::de::Visitor<'_> for Visitor {
        type Value = Vec<u8>;

        fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            formatter.write_str("a byte string")
        }

        fn visit_bytes<E: serde::de::Error>(self, bytes: &[u8]) -> std::result::Result<Vec<u8>, E> {
            Ok(bytes.to_vec())
        }

        fn visit_byte_buf<E: serde::de::Error>(
            self,
            bytes: Vec<u8>,
        ) -> std::result::Result<Vec<u8>, E> {
            Ok(bytes)
        }
    }
    deserializer.deserialize_byte_buf(Visitor)
}

/// What one entry recorded; the decoded view of an entry's content.
#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
#[serde(tag = "kind", content = "data")]
pub enum EntryContent {
    /// Chat message.
    Message(Message),
    /// Tool call (recorded for replay).
    ToolCall {
        /// Invoked tool.
        name: String,
        /// Canonical tool arguments.
        input: serde_json::Value,
        /// Identifier pairing this call with its result.
        call_id: ToolCallId,
    },
    /// Tool result.
    ToolResult {
        /// Whether the effect succeeded.
        ok: bool,
        /// Recorded output or error text.
        output: String,
        /// Identifier of the call this result answers.
        call_id: ToolCallId,
    },
    /// Session metadata event (compaction, branch selection).
    Meta {
        /// Metadata kind, for example [`COMPACTION_KIND`].
        kind: String,
        /// Kind-specific payload.
        detail: serde_json::Value,
    },
}

impl EntryContent {
    /// The record kind of this content.
    #[must_use]
    pub fn kind(&self) -> EntryKind {
        EntryContentRef::from(self).kind()
    }

    /// Decode content from its canonical CBOR bytes.
    ///
    /// # Errors
    ///
    /// Returns [`SmithError::Session`] when `bytes` are not entry content.
    pub fn decode(bytes: &[u8]) -> Result<Self> {
        ciborium::de::from_reader(bytes).map_err(|e| SmithError::Session {
            code: "ENTRY_DECODE".to_string(),
            message: e.to_string(),
        })
    }

    /// Append the canonical CBOR of this content to `out`.
    fn encode_into(&self, out: &mut Vec<u8>) -> Result<()> {
        EntryContentRef::from(self).encode_into(out)
    }
}

/// Borrowed [`EntryContent`]: what recording needs without owning a copy.
///
/// The serialize shape is the same as [`EntryContent`], so both produce the
/// same canonical bytes; this is the one place content is canonicalized.
#[derive(Clone, Copy, Debug, Serialize)]
#[serde(tag = "kind", content = "data")]
pub enum EntryContentRef<'a> {
    /// Chat message.
    Message(&'a Message),
    /// Tool call.
    ToolCall {
        /// Invoked tool.
        name: &'a str,
        /// Tool arguments.
        input: &'a serde_json::Value,
        /// Identifier pairing this call with its result.
        call_id: ToolCallId,
    },
    /// Tool result.
    ToolResult {
        /// Whether the effect succeeded.
        ok: bool,
        /// Output or error text.
        output: &'a str,
        /// Identifier of the call this result answers.
        call_id: ToolCallId,
    },
    /// Session metadata event.
    Meta {
        /// Metadata kind.
        kind: &'a str,
        /// Kind-specific payload.
        detail: &'a serde_json::Value,
    },
}

impl<'a> From<&'a EntryContent> for EntryContentRef<'a> {
    fn from(content: &'a EntryContent) -> Self {
        match content {
            EntryContent::Message(message) => Self::Message(message),
            EntryContent::ToolCall {
                name,
                input,
                call_id,
            } => Self::ToolCall {
                name,
                input,
                call_id: *call_id,
            },
            EntryContent::ToolResult {
                ok,
                output,
                call_id,
            } => Self::ToolResult {
                ok: *ok,
                output,
                call_id: *call_id,
            },
            EntryContent::Meta { kind, detail } => Self::Meta { kind, detail },
        }
    }
}

impl EntryContentRef<'_> {
    /// The record kind of this content.
    #[must_use]
    pub fn kind(&self) -> EntryKind {
        match *self {
            Self::Message(message) => EntryKind::Message {
                id: message.id,
                role: message.role,
            },
            Self::ToolCall { call_id, .. } => EntryKind::ToolCall { call_id },
            Self::ToolResult { call_id, ok, .. } => EntryKind::ToolResult { call_id, ok },
            Self::Meta { kind, .. } => EntryKind::Meta {
                compaction: kind == COMPACTION_KIND,
            },
        }
    }

    /// Append the canonical CBOR of this content to `out`; JSON values are
    /// re-ordered into canonical key order only when they are not already.
    fn encode_into(&self, out: &mut Vec<u8>) -> Result<()> {
        let canonical_value: serde_json::Value;
        let canonical_message: Message;
        // The view borrows from `self` or from a local canonical copy.
        let view: EntryContentRef<'_> = match *self {
            Self::ToolCall {
                name,
                input,
                call_id,
            } if !is_canonical(input) => {
                canonical_value = canonical_json(input);
                EntryContentRef::ToolCall {
                    name,
                    input: &canonical_value,
                    call_id,
                }
            }
            Self::Meta { kind, detail } if !is_canonical(detail) => {
                canonical_value = canonical_json(detail);
                EntryContentRef::Meta {
                    kind,
                    detail: &canonical_value,
                }
            }
            Self::Message(message) if !message_is_canonical(message) => {
                canonical_message = canonical_message_of(message);
                EntryContentRef::Message(&canonical_message)
            }
            other => other,
        };
        ciborium::ser::into_writer(&view, out).map_err(|e| SmithError::Session {
            code: "ENTRY_ENCODE".to_string(),
            message: e.to_string(),
        })
    }
}

fn message_is_canonical(message: &Message) -> bool {
    message.blocks.iter().all(|block| match block {
        ContentBlock::ToolUse { input, .. } => is_canonical(input),
        _ => true,
    })
}

fn canonical_message_of(message: &Message) -> Message {
    let mut message = message.clone();
    for block in &mut message.blocks {
        if let ContentBlock::ToolUse { input, .. } = block {
            *input = canonical_json(input);
        }
    }
    message
}

/// Wall-clock milliseconds since the Unix epoch, saturating on overflow.
fn now_ms() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |elapsed| {
            i64::try_from(elapsed.as_millis()).unwrap_or(i64::MAX)
        })
}

/// One root-to-leaf chain of a session.
///
/// Only its [`Session`] mutates a branch, so every record's span points into
/// that session's arena.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Branch {
    id: BranchId,
    entries: Vec<Entry>,
    selected_idx: usize,
}

impl Branch {
    const fn with_id(id: BranchId) -> Self {
        Self {
            id,
            entries: Vec::new(),
            selected_idx: 0,
        }
    }

    /// Branch identity.
    #[must_use]
    pub const fn id(&self) -> BranchId {
        self.id
    }

    /// Entry records ordered root to leaf.
    #[must_use]
    pub fn entries(&self) -> &[Entry] {
        &self.entries
    }

    /// Append `entry` and select it.
    fn append(&mut self, entry: Entry) -> &Entry {
        self.selected_idx = self.entries.len();
        self.entries.push(entry);
        &self.entries[self.selected_idx]
    }

    /// The currently selected entry, if the branch has any.
    #[must_use]
    pub fn selected_entry(&self) -> Option<&Entry> {
        self.entries.get(self.selected_idx)
    }
}

/// A session is a persistent, branching message/tool history.
///
/// Fields stay private so the invariants hold: at least one branch, a valid
/// active index, every span inside the arena, and a compaction boundary
/// that names a recorded compaction entry.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Session {
    id: SessionId,
    /// All branches (default branch is index 0).
    branches: Vec<Branch>,
    /// Canonical CBOR content of every entry, addressed by [`Span`].
    payloads: Vec<u8>,
    active_branch_idx: usize,
    compaction_boundary: Option<EntryId>,
}

impl Default for Session {
    fn default() -> Self {
        Self::new()
    }
}

impl Session {
    /// A new session with one empty branch.
    #[must_use]
    pub fn new() -> Self {
        Self {
            id: SessionId::new(),
            branches: vec![Branch::with_id(BranchId::new())],
            payloads: Vec::new(),
            active_branch_idx: 0,
            compaction_boundary: None,
        }
    }

    /// Session identifier.
    #[must_use]
    pub const fn id(&self) -> SessionId {
        self.id
    }

    /// All branches; the active one is [`Session::active_branch`].
    #[must_use]
    pub fn branches(&self) -> &[Branch] {
        &self.branches
    }

    /// The latest compaction boundary entry, if any.
    #[must_use]
    pub const fn compaction_boundary(&self) -> Option<EntryId> {
        self.compaction_boundary
    }

    /// The selected branch.
    #[must_use]
    pub fn active_branch(&self) -> &Branch {
        &self.branches[self.active_branch_idx]
    }

    /// Identifier of the selected entry, if any.
    #[must_use]
    pub fn selected_id(&self) -> Option<EntryId> {
        self.active_branch().selected_entry().map(|e| e.id)
    }

    /// Encode `content` once into the arena and append its record to the
    /// selected branch, selecting it. A compaction entry becomes the
    /// compaction boundary.
    ///
    /// # Errors
    ///
    /// Returns [`SmithError::Session`] when the content cannot be encoded or
    /// the arena is full; the session is unchanged then.
    pub fn append<'a>(&mut self, content: impl Into<EntryContentRef<'a>>) -> Result<&Entry> {
        self.append_with(content, |_, _| Ok(()))
    }

    /// [`Session::append`], running `persist` on the new record and its
    /// content bytes before the entry becomes part of the session.
    ///
    /// # Errors
    ///
    /// As [`Session::append`], or the error of `persist`; the session is
    /// unchanged then.
    pub fn append_with<'a>(
        &mut self,
        content: impl Into<EntryContentRef<'a>>,
        persist: impl FnOnce(&Entry, &[u8]) -> Result<()>,
    ) -> Result<&Entry> {
        let content = content.into();
        let parent = self.selected_id();
        let offset = self.payloads.len();
        let staged = content
            .encode_into(&mut self.payloads)
            .and_then(|()| Span::at(offset, self.payloads.len() - offset))
            .map(|span| Entry {
                id: EntryId::new(),
                parent,
                timestamp_ms: now_ms(),
                kind: content.kind(),
                content: span,
            })
            .and_then(|entry| persist(&entry, &self.payloads[offset..]).map(|()| entry));
        let entry = match staged {
            Ok(entry) => entry,
            Err(error) => {
                self.payloads.truncate(offset);
                return Err(error);
            }
        };
        if entry.is_compaction() {
            self.compaction_boundary = Some(entry.id);
        }
        Ok(self.branches[self.active_branch_idx].append(entry))
    }

    /// The canonical CBOR content bytes of `entry`; empty when `entry` does
    /// not belong to this session.
    #[must_use]
    pub fn bytes(&self, entry: &Entry) -> &[u8] {
        self.payloads.get(entry.content.range()).unwrap_or_default()
    }

    /// Decode the content of `entry` into the caller's scope, verifying it
    /// against the record's [`EntryKind`], which traversal trusts unchecked.
    ///
    /// # Errors
    ///
    /// Returns [`SmithError::Session`] with code `ENTRY_DECODE` when the
    /// bytes do not decode, or `ENTRY_KIND_MISMATCH` when they decode to
    /// content of another kind than the record claims.
    pub fn content(&self, entry: &Entry) -> Result<EntryContent> {
        let content = EntryContent::decode(self.bytes(entry))?;
        let found = content.kind();
        if found != entry.kind {
            return Err(SmithError::Session {
                code: "ENTRY_KIND_MISMATCH".to_string(),
                message: format!(
                    "entry {} records {:?} but holds {found:?}",
                    entry.id, entry.kind
                ),
            });
        }
        Ok(content)
    }

    /// The root-to-`at` entry chain, or `None` when `at` is unknown or its
    /// ancestry is broken.
    ///
    /// A branch stores its root-to-leaf chain in order, so the chain is the
    /// branch prefix ending at `at`; each parent link is verified against the
    /// preceding entry and no lookup table is built.
    fn ancestry(&self, at: EntryId) -> Option<&[Entry]> {
        let (branch, idx) = self.branches.iter().find_map(|branch| {
            let idx = branch.entries.iter().position(|entry| entry.id == at)?;
            Some((branch, idx))
        })?;
        let chain = &branch.entries[..=idx];
        let rooted = chain.first().is_some_and(|root| root.parent.is_none());
        let linked = chain
            .windows(2)
            .all(|pair| pair[1].parent == Some(pair[0].id));
        (rooted && linked).then_some(chain)
    }

    /// Fork the history ending at `at` into a new session.
    ///
    /// Entry IDs, order, and timestamps are preserved, `at` becomes the
    /// selected leaf, and the compaction boundary carries over when it lies
    /// on the retained path. The new session gets fresh session and branch
    /// identity, because it is a different durable history. The retained
    /// records are copied with the one arena prefix their spans cover, so
    /// spans stay valid unchanged.
    ///
    /// # Errors
    ///
    /// Returns [`SmithError::Session`] when `at` is not part of this session.
    pub fn fork_at(&self, at: EntryId) -> Result<Self> {
        let entries = self.ancestry(at).ok_or_else(|| SmithError::Session {
            code: "ENTRY_NOT_FOUND".to_string(),
            message: format!("entry {at} is not part of session {}", self.id),
        })?;
        let boundary = self
            .compaction_boundary
            .filter(|id| entries.iter().any(|entry| entry.id == *id));
        let end = entries
            .iter()
            .map(|entry| entry.content.end())
            .max()
            .unwrap_or(0);
        Ok(Self {
            id: SessionId::new(),
            branches: vec![Branch {
                id: BranchId::new(),
                entries: entries.to_vec(),
                selected_idx: entries.len().saturating_sub(1),
            }],
            payloads: self.payloads.get(..end).unwrap_or_default().to_vec(),
            active_branch_idx: 0,
            compaction_boundary: boundary,
        })
    }

    /// The durable identity of this session's active branch.
    #[must_use]
    pub fn header(&self) -> SessionHeader {
        SessionHeader {
            session_id: self.id,
            branch_id: self.active_branch().id,
        }
    }

    /// Restore in-memory state from decoded frames.
    ///
    /// A header frame restores the session and branch identity, so state
    /// rebuilt from the same file is identical on every load. Stored parent
    /// links define the branches: each entry with no children is a leaf, and
    /// its root-to-leaf chain is one branch. The last stored entry is the
    /// selected leaf of the active branch. Unknown frames are skipped here,
    /// and compaction entries restore the compaction boundary.
    ///
    /// Content bytes are copied into the arena, never decoded; every
    /// collection is sized up front, so a linear history restores with a
    /// constant number of allocations.
    ///
    /// # Errors
    ///
    /// Returns [`SmithError::Session`] when the content exceeds the arena cap.
    pub fn from_frames(frames: &[Frame]) -> Result<Self> {
        let mut header = None;
        let (mut count, mut total) = (0, 0);
        for frame in frames {
            match frame {
                Frame::Header { header: found } => header = Some(*found),
                Frame::Known { entry } => {
                    count += 1;
                    total += entry.content.len();
                }
                Frame::Unknown { .. } => {}
            }
        }
        Span::at(0, total)?;
        let mut payloads = Vec::with_capacity(total);
        let mut stored = Vec::with_capacity(count);
        let mut boundary = None;
        for frame in frames {
            let Frame::Known { entry } = frame else {
                continue;
            };
            let record = Entry {
                id: entry.id,
                parent: entry.parent,
                timestamp_ms: entry.timestamp_ms,
                kind: entry.kind,
                content: Span::at(payloads.len(), entry.content.len())?,
            };
            payloads.extend_from_slice(&entry.content);
            if record.is_compaction() {
                boundary = Some(record.id);
            }
            stored.push(record);
        }
        let header = header.unwrap_or_else(|| SessionHeader {
            session_id: SessionId::new(),
            branch_id: BranchId::new(),
        });
        let (branches, active_branch_idx) = restore_branches(&stored, header.branch_id);
        Ok(Self {
            id: header.session_id,
            branches,
            payloads,
            active_branch_idx,
            compaction_boundary: boundary,
        })
    }
}

/// Rebuild branches from stored records and their parent links, returning
/// them with the index of the active branch.
fn restore_branches(stored: &[Entry], active_id: BranchId) -> (Vec<Branch>, usize) {
    let Some(last) = stored.last() else {
        return (vec![Branch::with_id(active_id)], 0);
    };
    // Entry id to its index and whether any stored entry names it as parent.
    let mut nodes: HashMap<EntryId, (usize, bool)> = HashMap::with_capacity(stored.len());
    for (idx, entry) in stored.iter().enumerate() {
        nodes.insert(entry.id, (idx, false));
    }
    for parent in stored.iter().filter_map(|entry| entry.parent) {
        if let Some(node) = nodes.get_mut(&parent) {
            node.1 = true;
        }
    }
    let nodes = &nodes;
    let is_leaf = |entry: &&Entry| nodes.get(&entry.id).is_some_and(|node| !node.1);
    // Root-ward walk; a missing parent ends the recoverable chain, and the
    // step cap ends a corrupt cycle.
    let chain = |leaf: &Entry| {
        let mut cursor = Some(leaf.id);
        std::iter::from_fn(move || {
            let (idx, _) = nodes.get(&cursor?)?;
            let entry = stored.get(*idx)?;
            cursor = entry.parent;
            Some(*entry)
        })
        .take(stored.len())
    };
    let mut branches = Vec::with_capacity(stored.iter().filter(is_leaf).count());
    let mut active_branch_idx = 0;
    for leaf in stored.iter().filter(is_leaf) {
        let mut entries = Vec::with_capacity(chain(leaf).count());
        entries.extend(chain(leaf));
        entries.reverse();
        let is_active = leaf.id == last.id;
        if is_active {
            active_branch_idx = branches.len();
        }
        branches.push(Branch {
            // The selected branch owns the persisted identity; sibling
            // branches derive theirs from their leaf, which keeps every
            // reload of the same file identical.
            id: if is_active {
                active_id
            } else {
                BranchId::from_u128(leaf.id.as_u128())
            },
            selected_idx: entries.len().saturating_sub(1),
            entries,
        });
    }
    (branches, active_branch_idx)
}

#[cfg(test)]
#[expect(
    clippy::unwrap_used,
    clippy::panic,
    reason = "tests may panic on invariant violations"
)]
mod tests {
    use super::*;

    fn message(text: &str) -> EntryContent {
        EntryContent::Message(Message::with_text(Role::User, text))
    }

    fn compaction(detail: serde_json::Value) -> EntryContent {
        EntryContent::Meta {
            kind: COMPACTION_KIND.to_string(),
            detail,
        }
    }

    fn text_of(session: &Session, entry: &Entry) -> String {
        match session.content(entry).unwrap() {
            EntryContent::Message(m) => format!("{}", m.blocks.len()),
            EntryContent::Meta { kind, .. } => kind,
            other => panic!("unexpected content {other:?}"),
        }
    }

    /// The frame of every record in `entries`, content copied from `session`.
    fn frames_of(session: &Session, entries: &[Entry]) -> Vec<Frame> {
        let mut frames = vec![Frame::Header {
            header: session.header(),
        }];
        frames.extend(entries.iter().map(|entry| Frame::Known {
            entry: EntryFrame {
                id: entry.id,
                parent: entry.parent,
                timestamp_ms: entry.timestamp_ms,
                kind: entry.kind,
                content: session.bytes(entry).to_vec(),
            },
        }));
        frames
    }

    fn frames(session: &Session) -> Vec<Frame> {
        frames_of(session, &session.active_branch().entries)
    }

    /// Frames of entries built outside any session, under a fresh header;
    /// the only way to store parent links a session would not append.
    fn stored(entries: &[EntryFrame]) -> (SessionHeader, Vec<Frame>) {
        let header = SessionHeader {
            session_id: SessionId::new(),
            branch_id: BranchId::new(),
        };
        let mut frames = vec![Frame::Header { header }];
        frames.extend(entries.iter().cloned().map(|entry| Frame::Known { entry }));
        (header, frames)
    }

    fn ids(branch: &Branch) -> Vec<EntryId> {
        branch.entries.iter().map(|entry| entry.id).collect()
    }

    #[test]
    fn sibling_entries_restore_as_separate_branches() {
        // Stored parent links record two children of one root.
        let root = EntryFrame::new(None, &message("root")).unwrap();
        let first = EntryFrame::new(Some(root.id), &message("first")).unwrap();
        let second = EntryFrame::new(Some(root.id), &message("second")).unwrap();
        let (header, stored) = stored(&[root.clone(), first.clone(), second.clone()]);
        let session = Session::from_frames(&stored).unwrap();
        assert_eq!(session.branches.len(), 2);
        // The last stored entry selects the active branch and keeps the
        // persisted branch identity.
        assert_eq!(session.selected_id(), Some(second.id));
        assert_eq!(session.active_branch().id, header.branch_id);
        assert_eq!(ids(session.active_branch()), vec![root.id, second.id]);
        let sibling = &session.branches[1 - session.active_branch_idx];
        assert_eq!(ids(sibling), vec![root.id, first.id]);
        // Shared-prefix records point at the same content bytes.
        assert_eq!(
            session.content(&sibling.entries[0]).unwrap(),
            EntryContent::decode(&root.content).unwrap()
        );
        // Reconstruction is identical on a second pass, sibling ids included.
        let again = Session::from_frames(&stored).unwrap();
        assert_eq!(again, session);
    }

    #[test]
    fn appended_content_decodes_from_the_arena_and_records_its_kind() {
        let mut session = Session::new();
        let content = EntryContent::ToolResult {
            ok: false,
            output: "boom".to_string(),
            call_id: ToolCallId::new(),
        };
        let entry = *session.append(&content).unwrap();
        assert_eq!(entry.kind, content.kind());
        assert_eq!(session.content(&entry).unwrap(), content);
        assert_eq!(session.bytes(&entry).len(), entry.content.len as usize);
        // A record from elsewhere reads as empty bytes and fails to decode.
        let foreign = Entry {
            content: Span {
                offset: 1 << 20,
                len: 4,
            },
            ..entry
        };
        assert!(session.bytes(&foreign).is_empty());
        assert_eq!(
            session.content(&foreign).unwrap_err().code(),
            "ENTRY_DECODE"
        );
    }

    #[test]
    fn failed_persist_leaves_the_session_unchanged() {
        let mut session = Session::new();
        session.append(&message("kept")).unwrap();
        let before = session.clone();
        let err = session
            .append_with(
                &EntryContent::Meta {
                    kind: COMPACTION_KIND.to_string(),
                    detail: serde_json::json!({}),
                },
                |_, bytes| {
                    assert!(!bytes.is_empty());
                    Err(SmithError::Cancelled)
                },
            )
            .unwrap_err();
        assert_eq!(err.code(), "CANCELLED");
        assert_eq!(session, before);
    }

    #[test]
    fn fork_preserves_ids_order_leaf_and_compaction_boundary() {
        let mut session = Session::new();
        let root = *session.append(&message("root")).unwrap();
        let boundary = *session
            .append(&compaction(serde_json::json!({"summarized": 1})))
            .unwrap();
        let after = *session.append(&message("after")).unwrap();

        let forked = session.fork_at(boundary.id).unwrap();
        assert_ne!(forked.id, session.id);
        assert_eq!(ids(forked.active_branch()), vec![root.id, boundary.id]);
        assert_eq!(forked.selected_id(), Some(boundary.id));
        assert_eq!(forked.compaction_boundary, Some(boundary.id));
        assert_eq!(
            forked
                .active_branch()
                .entries
                .iter()
                .map(|entry| text_of(&forked, entry))
                .collect::<Vec<_>>(),
            vec!["1".to_string(), COMPACTION_KIND.to_string()]
        );
        // The fork carries only the arena prefix its records cover.
        assert_eq!(forked.bytes(&after), &[] as &[u8]);

        // Forking before the boundary drops it, because it is not retained.
        let earlier = session.fork_at(root.id).unwrap();
        assert_eq!(earlier.compaction_boundary, None);
        assert_eq!(earlier.active_branch().entries.len(), 1);

        // The source session is untouched.
        assert_eq!(session.active_branch().entries.len(), 3);
        assert_eq!(session.selected_id(), Some(after.id));
    }

    #[test]
    fn tool_use_json_inside_messages_encodes_identically() {
        let call_id = ToolCallId::new();
        let block = |input: serde_json::Value| {
            let mut message = Message::with_text(Role::Assistant, "call");
            message.add_block(ContentBlock::ToolUse {
                name: "write".to_string(),
                input,
                call_id,
            });
            message
        };
        let first = block(serde_json::json!({"path": "p", "content": "c"}));
        let mut second = block(serde_json::json!({"content": "c", "path": "p"}));
        // Only key order may differ, so encodings must match.
        second.id = first.id;
        second.timestamp_ms = first.timestamp_ms;
        let mut session = Session::new();
        let a = *session.append(&EntryContent::Message(first)).unwrap();
        let b = *session.append(&EntryContent::Message(second)).unwrap();
        assert_eq!(session.bytes(&a), session.bytes(&b));
    }

    #[test]
    fn fork_rejects_unknown_entries() {
        let session = Session::new();
        let err = session.fork_at(EntryId::new()).unwrap_err();
        assert_eq!(err.code(), "ENTRY_NOT_FOUND");
    }

    #[test]
    fn fork_rejects_broken_ancestry_and_forks_sibling_branches() {
        // A stored chain whose root parent is missing restores, but cannot fork.
        let orphan = EntryFrame::new(Some(EntryId::new()), &message("orphan")).unwrap();
        let child = EntryFrame::new(Some(orphan.id), &message("child")).unwrap();
        let broken = Session::from_frames(&stored(&[orphan, child.clone()]).1).unwrap();
        assert_eq!(broken.selected_id(), Some(child.id));
        let err = broken.fork_at(child.id).unwrap_err();
        assert_eq!(err.code(), "ENTRY_NOT_FOUND");

        // A leaf on an inactive sibling branch forks along its own chain.
        let root = EntryFrame::new(None, &message("root")).unwrap();
        let first = EntryFrame::new(Some(root.id), &message("first")).unwrap();
        let second = EntryFrame::new(Some(root.id), &message("second")).unwrap();
        let siblings =
            Session::from_frames(&stored(&[root.clone(), first.clone(), second]).1).unwrap();
        let forked = siblings.fork_at(first.id).unwrap();
        assert_eq!(ids(forked.active_branch()), vec![root.id, first.id]);
        assert_eq!(text_of(&forked, &forked.active_branch().entries[1]), "1");
    }

    #[test]
    fn loading_copies_content_bytes_without_decoding_them() {
        let mut scratch = Session::new();
        let root = *scratch.append(&message("root")).unwrap();
        let mut stored = frames(&scratch);
        if let Frame::Known { entry } = &mut stored[1] {
            entry.content = vec![0xFF, 0x00];
        }
        let restored = Session::from_frames(&stored).unwrap();
        let record = restored.active_branch().entries[0];
        assert_eq!(record.kind, root.kind);
        assert_eq!(restored.bytes(&record), &[0xFF, 0x00]);
        assert_eq!(
            restored.content(&record).unwrap_err().code(),
            "ENTRY_DECODE"
        );
    }

    #[test]
    fn content_of_a_record_whose_kind_lies_is_a_kind_mismatch() {
        let mut scratch = Session::new();
        scratch.append(&message("root")).unwrap();
        let mut stored = frames(&scratch);
        let lie = EntryKind::ToolCall {
            call_id: ToolCallId::new(),
        };
        if let Frame::Known { entry } = &mut stored[1] {
            entry.kind = lie;
        }
        // Loading trusts the record; only a decode can tell.
        let restored = Session::from_frames(&stored).unwrap();
        let record = restored.active_branch().entries[0];
        assert_eq!(record.kind, lie);
        assert_eq!(
            restored.content(&record).unwrap_err().code(),
            "ENTRY_KIND_MISMATCH"
        );
    }

    #[test]
    fn compaction_boundary_survives_a_frame_round_trip() {
        let mut session = Session::new();
        session.append(&message("root")).unwrap();
        let boundary = *session.append(&compaction(serde_json::json!({}))).unwrap();
        let restored = Session::from_frames(&frames(&session)).unwrap();
        assert_eq!(restored.compaction_boundary, Some(boundary.id));
        assert_eq!(restored.id, session.id);
        assert_eq!(restored.active_branch().id, session.header().branch_id);
        assert_eq!(restored, session);
    }
}

/// Laws of canonicalization, span arithmetic, and content encoding
/// (`SMH-SPEC-SPEC0001`, Sessions).
///
/// `serde_json` preserves insertion order but compares maps order-blind, so
/// key-order laws compare serialized bytes, never values.
#[cfg(test)]
pub mod properties {
    use super::*;
    use proptest::collection::vec;
    use proptest::prelude::*;
    use serde_json::Value;

    /// Arbitrary JSON, at most three levels deep and 24 nodes large.
    fn json() -> impl Strategy<Value = Value> {
        let leaf = prop_oneof![
            Just(Value::Null),
            any::<bool>().prop_map(Value::from),
            any::<i64>().prop_map(Value::from),
            any::<u64>().prop_map(Value::from),
            (proptest::num::f64::NORMAL | proptest::num::f64::ZERO).prop_map(Value::from),
            "\\PC{0,8}".prop_map(Value::from),
        ];
        leaf.prop_recursive(3, 24, 4, |inner| {
            prop_oneof![
                vec(inner.clone(), 0..4).prop_map(Value::Array),
                vec(("\\PC{0,4}", inner), 0..5)
                    .prop_map(|entries| Value::Object(entries.into_iter().collect())),
            ]
        })
    }

    /// `value` with the keys of every object reordered as `seed` dictates.
    fn reordered(value: &Value, seed: u64) -> Value {
        match value {
            Value::Object(map) => {
                let mut entries: Vec<(String, Value)> = map
                    .iter()
                    .map(|(key, value)| (key.clone(), reordered(value, seed.rotate_left(7))))
                    .collect();
                if !entries.is_empty() {
                    let by = usize::try_from(seed % entries.len() as u64).unwrap_or(0);
                    entries.rotate_left(by);
                }
                if seed & 1 == 1 {
                    entries.reverse();
                }
                Value::Object(entries.into_iter().collect())
            }
            Value::Array(items) => Value::Array(
                items
                    .iter()
                    .map(|item| reordered(item, seed.rotate_left(7)))
                    .collect(),
            ),
            other => other.clone(),
        }
    }

    fn json_bytes(value: &Value) -> std::result::Result<Vec<u8>, TestCaseError> {
        serde_json::to_vec(value).map_err(|e| TestCaseError::fail(e.to_string()))
    }

    fn cbor<T: Serialize>(value: &T) -> std::result::Result<Vec<u8>, TestCaseError> {
        let mut out = Vec::new();
        ciborium::ser::into_writer(value, &mut out)
            .map(|()| out)
            .map_err(|e| TestCaseError::fail(e.to_string()))
    }

    /// Any message role.
    pub fn role() -> impl Strategy<Value = Role> {
        prop_oneof![
            Just(Role::User),
            Just(Role::Assistant),
            Just(Role::System),
            Just(Role::Tool),
        ]
    }

    fn block() -> impl Strategy<Value = ContentBlock> {
        prop_oneof![
            "\\PC{0,16}".prop_map(ContentBlock::Text),
            "\\PC{0,16}".prop_map(ContentBlock::Thinking),
            ("\\PC{0,8}", json(), any::<u128>()).prop_map(|(name, input, id)| {
                ContentBlock::tool_use(name, input, ToolCallId::from_u128(id))
            }),
            (any::<u128>(), any::<bool>(), "\\PC{0,16}").prop_map(|(id, ok, output)| {
                ContentBlock::tool_result(ToolCallId::from_u128(id), ok, output)
            }),
        ]
    }

    fn content() -> impl Strategy<Value = EntryContent> {
        prop_oneof![
            (any::<u128>(), role(), vec(block(), 0..4), any::<i64>()).prop_map(
                |(id, role, blocks, timestamp_ms)| EntryContent::Message(Message {
                    id: MessageId::from_u128(id),
                    role,
                    blocks,
                    timestamp_ms,
                })
            ),
            ("\\PC{0,8}", json(), any::<u128>()).prop_map(|(name, input, id)| {
                EntryContent::ToolCall {
                    name,
                    input,
                    call_id: ToolCallId::from_u128(id),
                }
            }),
            (any::<bool>(), "\\PC{0,16}", any::<u128>()).prop_map(|(ok, output, id)| {
                EntryContent::ToolResult {
                    ok,
                    output,
                    call_id: ToolCallId::from_u128(id),
                }
            }),
            (
                prop_oneof![
                    Just(COMPACTION_KIND.to_string()),
                    Just(PROVIDER_HTTP_KIND.to_string()),
                    "\\PC{0,8}",
                ],
                json()
            )
                .prop_map(|(kind, detail)| EntryContent::Meta { kind, detail }),
        ]
    }

    /// `content` with every JSON value in canonical key order.
    fn canonical_content(content: &EntryContent) -> EntryContent {
        match content {
            EntryContent::Message(message) => EntryContent::Message(canonical_message_of(message)),
            EntryContent::ToolCall {
                name,
                input,
                call_id,
            } => EntryContent::ToolCall {
                name: name.clone(),
                input: canonical_json(input),
                call_id: *call_id,
            },
            EntryContent::ToolResult { .. } => content.clone(),
            EntryContent::Meta { kind, detail } => EntryContent::Meta {
                kind: kind.clone(),
                detail: canonical_json(detail),
            },
        }
    }

    /// Offsets and lengths around the 4 GiB arena cap, plus anything.
    fn arena_size() -> impl Strategy<Value = usize> {
        let cap = u32::MAX as usize;
        prop_oneof![
            0..=cap,
            cap - 64..=cap + 64,
            (0..=64_usize).prop_map(move |below| cap - below),
            any::<usize>(),
        ]
    }

    proptest! {
        #[test]
        fn canonical_json_is_idempotent_and_keeps_the_value(value in json()) {
            let once = canonical_json(&value);
            prop_assert_eq!(json_bytes(&canonical_json(&once))?, json_bytes(&once)?);
            prop_assert!(is_canonical(&once));
            prop_assert_eq!(once, value);
        }

        #[test]
        fn canonical_json_ignores_key_order(value in json(), seed in any::<u64>()) {
            let shuffled = reordered(&value, seed);
            prop_assert_eq!(
                json_bytes(&canonical_json(&shuffled))?,
                json_bytes(&canonical_json(&value))?
            );
        }

        #[test]
        fn is_canonical_holds_exactly_when_canonicalizing_changes_nothing(value in json()) {
            prop_assert_eq!(
                is_canonical(&value),
                json_bytes(&canonical_json(&value))? == json_bytes(&value)?
            );
        }

        #[test]
        fn span_at_is_exact_below_the_cap_and_an_error_above(
            offset in arena_size(),
            len in arena_size(),
        ) {
            let fits = offset
                .checked_add(len)
                .is_some_and(|end| u32::try_from(end).is_ok());
            match Span::at(offset, len) {
                Ok(span) => {
                    prop_assert!(fits);
                    prop_assert_eq!(span.offset as usize, offset);
                    prop_assert_eq!(span.len as usize, len);
                    prop_assert_eq!(span.end(), offset + len);
                    prop_assert_eq!(span.range().len(), len);
                }
                Err(error) => {
                    prop_assert!(!fits);
                    prop_assert_eq!(error.code(), "SESSION_ARENA_FULL");
                }
            }
        }

        #[test]
        fn content_ref_serializes_like_owned_content(content in content()) {
            prop_assert_eq!(cbor(&EntryContentRef::from(&content))?, cbor(&content)?);
        }

        #[test]
        fn owned_and_borrowed_encodings_are_canonical_and_decode_back(content in content()) {
            let owned = EntryFrame::new(None, &content)
                .map_err(|e| TestCaseError::fail(e.to_string()))?
                .content;
            let mut session = Session::new();
            let record = *session
                .append(EntryContentRef::from(&content))
                .map_err(|e| TestCaseError::fail(e.to_string()))?;
            prop_assert_eq!(session.bytes(&record), owned.as_slice());
            prop_assert_eq!(&owned, &cbor(&canonical_content(&content))?);
            prop_assert_eq!(
                EntryContent::decode(&owned).map_err(|e| TestCaseError::fail(e.to_string()))?,
                content
            );
        }
    }
}