//! Streaming MST construction stack — the algorithm-of-record for
//! STAR-lite-shaped MST construction.
//!
//! Specialized for sorted-key input: `(key, record_bytes)` pairs arrive in
//! strict lex order, and the stack produces a canonical atproto MST. No
//! rebalancing, no random updates, no general-purpose semantics.
//!
//! ## Algorithm
//!
//! The stack maintains one partial node per MST layer. When a new key K at
//! layer D arrives:
//!
//! 1. For each layer L in `0..D`, finalize whatever's accumulated: encode
//!    its slots as a single MST node (always — no trivial-pass-through; per
//!    AT-Repo §2.5.4 a layer with just a subtree slot becomes the
//!    intermediate empty-entry node `{l: cid, e: []}`). Push the resulting
//!    subtree onto layer L+1.
//! 2. Push a leaf for K onto layer D.
//!
//! At [`MstStack::finish`], we fold the remaining layers bottom-up the same
//! way. Empty layers with a carry from below produce the bridge nodes
//! described above. Empty layers with no carry produce nothing — that's how
//! "top pruning" (the spec's term for not adding empty wrappers above the
//! highest-layer key) and "leaf-position pruning" both fall out naturally.
//!
//! For the empty-repo case (no inserts) we emit the single empty node
//! `{l: null, e: []}` per §2.5.4.
//!
//! ## What the stack owns vs. what backends own
//!
//! The stack owns *everything* that affects the wire format: ordering checks,
//! layer promotion, the entry cap, MST node encoding, prefix elision. There
//! is exactly one source of truth for these — verifiers and converters cannot
//! desync, because they all run the same code.
//!
//! Backends own *only* what to do with bytes the stack produces: persist
//! them, hash and discard them, etc. See [`MstBackend`].

use star_repo::Cid;

use star_repo::compute_cid_dag_cbor;
use star_repo::error::MstError;
use star_repo::mst::layer_of;
use star_repo::mst::node::{MstEntry, MstNode, encode_node, shared_prefix_len};

use crate::error::Result;

/// Default cap on entries in any single MST node. Atproto MSTs naturally have
/// a few entries per node (fanout 4); 200 is a generous safety margin
/// matching the per-commit operation cap from §4.2.1.
pub const DEFAULT_MAX_ENTRIES_PER_NODE: usize = 200;

#[derive(Debug, Clone)]
pub struct BuilderOptions {
    pub max_entries_per_node: usize,
}

impl Default for BuilderOptions {
    fn default() -> Self {
        Self {
            max_entries_per_node: DEFAULT_MAX_ENTRIES_PER_NODE,
        }
    }
}

/// Connects the [`MstStack`] state machine to a per-format storage strategy.
///
/// Implementations choose what to do with the bytes the stack produces — some
/// (verify) hash and discard, some (CAR conversion) persist into a byte log.
/// Backends never see ordering, layer-promotion, encoding, or cap-enforcement
/// logic; those live in the stack, so two backends running on the same input
/// cannot disagree about the wire format.
pub trait MstBackend {
    /// Per-leaf data the backend tracks. `()` if it doesn't need any.
    type LeafRef;
    /// Per-subtree handle the backend produces when a layer finalizes. The
    /// CID is required (via [`Self::slot_cid`]) so the stack can name this
    /// slot in the parent node above it.
    type Slot;

    /// Take ownership of an inserted record's bytes. The stack has already
    /// computed `record_cid`; the backend may persist the bytes to a byte
    /// log, hash them out-of-band, or simply ignore them.
    fn ingest_leaf(&mut self, record_cid: Cid, record_bytes: &[u8]) -> Result<Self::LeafRef>;

    /// Take ownership of a finalized layer: a fully-encoded MST node (CID +
    /// bytes), plus its children. Produce a single slot representing the
    /// resulting subtree.
    ///
    /// `node_bytes` is the canonical DAG-CBOR encoding the stack just
    /// produced, and `node_cid` is its CID. Backends that persist bytes
    /// should use these directly rather than re-encoding.
    fn build_slot(
        &mut self,
        node_cid: Cid,
        node_bytes: &[u8],
        left: Option<Self::Slot>,
        entries: Vec<LevelEntry<Self::LeafRef, Self::Slot>>,
    ) -> Result<Self::Slot>;

    /// Read the root CID out of a slot. Called by the stack when assembling
    /// parent nodes (the parent needs each child's CID to encode its `l` and
    /// `t` fields).
    fn slot_cid(slot: &Self::Slot) -> Cid;
}

/// One leaf entry, plus the optional right subtree that follows it. Mirrors
/// the encoded MST node 1:1.
#[derive(Debug)]
pub struct LevelEntry<L, S> {
    pub key: Vec<u8>,
    pub leaf_cid: Cid,
    pub leaf_ref: L,
    pub right: Option<S>,
}

/// One layer of the stack: an optional left subtree, followed by leaves
/// (each with an optional right subtree).
struct Level<L, S> {
    left: Option<S>,
    entries: Vec<LevelEntry<L, S>>,
}

impl<L, S> Default for Level<L, S> {
    fn default() -> Self {
        Self {
            left: None,
            entries: Vec::new(),
        }
    }
}

impl<L, S> Level<L, S> {
    fn is_empty(&self) -> bool {
        self.left.is_none() && self.entries.is_empty()
    }

    /// Push a freshly-promoted subtree at the right end of this layer.
    /// Subtrees attach as `left` if no leaves have been seen yet, otherwise
    /// as `right` of the most recent leaf.
    fn push_subtree(&mut self, s: S) -> Result<()> {
        if self.entries.is_empty() {
            if self.left.is_some() {
                return Err(
                    MstError::Invariant("two consecutive subtree slots before any leaf").into(),
                );
            }
            self.left = Some(s);
        } else {
            let tail = self
                .entries
                .last_mut()
                .expect("entries non-empty by branch check");
            if tail.right.is_some() {
                return Err(
                    MstError::Invariant("two consecutive subtree slots after a leaf").into(),
                );
            }
            tail.right = Some(s);
        }
        Ok(())
    }
}

pub struct MstStack<B: MstBackend> {
    options: BuilderOptions,
    backend: B,
    layers: Vec<Level<B::LeafRef, B::Slot>>,
    last_key: Option<Vec<u8>>,
}

impl<B: MstBackend> MstStack<B> {
    pub fn new(backend: B, options: BuilderOptions) -> Self {
        Self {
            options,
            backend,
            layers: Vec::new(),
            last_key: None,
        }
    }

    /// Insert a `(key, record)` pair. Keys must arrive in strict lex order.
    pub fn insert(&mut self, key: &[u8], record: &[u8]) -> Result<()> {
        if key.is_empty() {
            return Err(MstError::EmptyKey.into());
        }
        if let Some(prev) = &self.last_key
            && key <= prev.as_slice()
        {
            return Err(MstError::KeyOrder {
                prev: String::from_utf8_lossy(prev).into_owned(),
                curr: String::from_utf8_lossy(key).into_owned(),
            }
            .into());
        }

        let layer = layer_of(key) as usize;
        while self.layers.len() <= layer {
            self.layers.push(Level::default());
        }

        // Promote: every layer below this key's layer is now frozen.
        for level_idx in 0..layer {
            let level = std::mem::take(&mut self.layers[level_idx]);
            if level.is_empty() {
                continue;
            }
            let promoted = encode_level(&mut self.backend, level, &self.options)?;
            self.layers[level_idx + 1].push_subtree(promoted)?;
        }

        // Eager cap check: enforce the same limit at accumulation time as
        // encode_level enforces at finalization. Without this, an adversary
        // submitting a repo with a long run of low-layer keys (cheap to
        // grind: ~75% of random keys are layer 0) could grow this layer's
        // entries vector unboundedly before the encode-time check fires,
        // OOMing the process. Living in the stack — not in each backend —
        // means verifiers and converters fail at *exactly* the same input
        // boundary, so an attacker can't ship a repo that one accepts and
        // another rejects.
        if self.layers[layer].entries.len() >= self.options.max_entries_per_node {
            return Err(MstError::NodeTooLarge {
                count: self.layers[layer].entries.len() + 1,
                limit: self.options.max_entries_per_node,
            }
            .into());
        }

        let record_cid = compute_cid_dag_cbor(record);
        let leaf_ref = self.backend.ingest_leaf(record_cid, record)?;
        self.layers[layer].entries.push(LevelEntry {
            key: key.to_vec(),
            leaf_cid: record_cid,
            leaf_ref,
            right: None,
        });
        self.last_key = Some(key.to_vec());
        Ok(())
    }

    /// Finalize the tree. Returns `(root_slot, backend)` so callers can read
    /// the root CID and continue using any resources the backend holds (e.g.
    /// the populated byte log a CAR emitter needs to copy from).
    ///
    /// If no inserts happened, returns the canonical empty MST node.
    pub fn finish(mut self) -> Result<(B::Slot, B)> {
        let mut carry: Option<B::Slot> = None;
        let depth = self.layers.len();
        for level_idx in 0..depth {
            let mut level = std::mem::take(&mut self.layers[level_idx]);
            if let Some(c) = carry.take() {
                // The carry attaches at the right end. push_subtree handles
                // both empty-layer (becomes left) and non-empty-layer
                // (attaches to last entry's right) cases, and rejects the
                // invariant violations.
                level.push_subtree(c)?;
            }
            if level.is_empty() {
                continue;
            }
            // Always encode — never trivial-collapse a single-subtree layer.
            // Per §2.5.4 a layer with just a subtree carry produces the
            // intermediate `{l: cid, e: []}` bridge node, *unless* it ends up
            // being the root. By construction this can't happen here: the
            // topmost stack layer always has at least one leaf (it's the
            // layer of the highest-layer key, which was inserted as a leaf
            // there). So if `level_idx` is below the top, encoding a bridge
            // is correct; if `level_idx` is the top, the encoded node has
            // entries.
            carry = Some(encode_level(&mut self.backend, level, &self.options)?);
        }

        let root = match carry {
            Some(r) => r,
            None => encode_empty_root(&mut self.backend)?,
        };
        Ok((root, self.backend))
    }
}

/// Encode a layer as a single MST node and hand the result to the backend.
fn encode_level<B: MstBackend>(
    backend: &mut B,
    level: Level<B::LeafRef, B::Slot>,
    options: &BuilderOptions,
) -> Result<B::Slot> {
    debug_assert!(!level.is_empty(), "encode_level called with empty level");

    // Backstop. The eager check in `MstStack::insert` should make this
    // unreachable in practice, but `encode_level` may be called with
    // synthesized layers during `finish` (when a carry from below attaches
    // to a layer that was empty), so we keep this as a safety net.
    if level.entries.len() > options.max_entries_per_node {
        return Err(MstError::NodeTooLarge {
            count: level.entries.len(),
            limit: options.max_entries_per_node,
        }
        .into());
    }

    let Level { left, entries } = level;

    let left_cid = left.as_ref().map(B::slot_cid);
    let mut mst_entries: Vec<MstEntry> = Vec::with_capacity(entries.len());
    let mut prev_key: Option<&[u8]> = None;
    for entry in &entries {
        let prefix_len = match prev_key {
            None => 0,
            Some(p) => shared_prefix_len(p, &entry.key),
        };
        let suffix = entry.key[prefix_len..].to_vec();
        mst_entries.push(MstEntry {
            prefix_len: prefix_len as u32,
            key_suffix: serde_bytes::ByteBuf::from(suffix),
            value: entry.leaf_cid,
            right: entry.right.as_ref().map(B::slot_cid),
        });
        prev_key = Some(&entry.key);
    }

    let node = MstNode {
        entries: mst_entries,
        left: left_cid,
    };
    let (cid, bytes) = encode_node(&node)?;
    backend.build_slot(cid, &bytes, left, entries)
}

fn encode_empty_root<B: MstBackend>(backend: &mut B) -> Result<B::Slot> {
    let node = MstNode {
        entries: Vec::new(),
        left: None,
    };
    let (cid, bytes) = encode_node(&node)?;
    backend.build_slot(cid, &bytes, None, Vec::new())
}