//! A state: a dependency-closed set of changes, applied to a tree graph and

//! per-node file graphs (spec §8.3), with materialization (spec §9.5, §11).

use crate::cbor::{self, Value};

use crate::graph::{AtomKey, DepOracle, FileGraph, Materialized, Vertex};

use crate::hash::KeyId;

use crate::hash::{ChangeHash, ObjectRef};

use crate::ids::NodeId;

use crate::objects::{

    self as o, AtomRef, Change, ChangeKind, Edit, Hlc, Mode, NodeKind, ObjectError, TreeOp,

};

use crate::tree::{TreeGraph, TreeView};

use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};

use std::sync::{Arc, Mutex};

/// Checkpoint payload version 2 (spec §11): file graphs by reference.

pub const CHECKPOINT_INDEX_VERSION: u64 = 2;

/// A file graph encoded in fewer bytes stays inline in a version 2 index.

pub const INLINE_GRAPH_LIMIT: usize = 2048;

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]

pub enum StateError {

    #[error("change {1} depends on {0}, which is not in the state")]

    MissingDep(ChangeHash, ChangeHash),

}

/// One `Replace` edit: `(change, old blob, new blob)`.

type ReplaceRecord = (ChangeHash, Option<ObjectRef>, Option<ObjectRef>);

/// What a state remembers about a change: everything but its operations,

/// which live in the graphs once applied (and in the store).

#[derive(Clone, Debug, PartialEq, Eq)]

pub struct ChangeMeta {

    pub change_id: crate::ids::ChangeId,

    pub deps: Vec<ChangeHash>,

    pub author: KeyId,

    pub hlc: Hlc,

    pub message: String,

    pub kind: ChangeKind,

}

impl ChangeMeta {

    /// In the flat layout of local indexes and caches ([`crate::flat`]).

    pub(crate) fn write_flat(&self, w: &mut crate::flat::Writer) {

        w.fixed(&self.change_id.0);

        w.count(self.deps.len());

        for d in &self.deps {

            w.object(d);

        }

        w.object(&self.author);

        w.u64(self.hlc.wall_ms);

        w.u32(self.hlc.logical);

        w.bytes(self.message.as_bytes());

        w.u8(self.kind.tag() as u8);

    }

    pub(crate) fn read_flat(r: &mut crate::flat::Reader) -> Result<ChangeMeta, ObjectError> {

        let change_id = crate::ids::ChangeId(r.fixed::<16>()?);

        let mut deps = Vec::new();

        for _ in 0..r.count()? {

            deps.push(r.object()?);

        }

        Ok(ChangeMeta {

            change_id,

            deps,

            author: r.object()?,

            hlc: Hlc {

                wall_ms: r.u64()?,

                logical: r.u32()?,

            },

            message: r.text()?,

            kind: ChangeKind::from_tag(r.u8()? as u64)?,

        })

    }

    pub fn of(c: &Change) -> ChangeMeta {

        ChangeMeta {

            change_id: c.change_id,

            deps: c.deps.clone(),

            author: c.author,

            hlc: c.hlc,

            message: c.message.clone(),

            kind: c.kind,

        }

    }

}

#[derive(Default)]

pub struct State {

    /// Shared between clones, like `files` and `deferred`: a state kept

    /// beside another that differs by a few changes shares the rest.

    changes: HashMap<ChangeHash, Arc<ChangeMeta>>,

    tree: TreeGraph,

    /// Shared between clones: a clone copies a graph only when an edit

    /// changes it (`Arc::make_mut`), so keeping several states that differ

    /// by a few changes costs little more than one.

    files: crate::fasthash::FastMap<NodeId, Arc<FileGraph>>,

    replaces: HashMap<NodeId, Vec<ReplaceRecord>>,

    /// Memoised transitive dependencies. A mutex, not a `RefCell`, so a

    /// loaded state can be shared across threads by a host that caches it.

    closure: Mutex<HashMap<ChangeHash, Arc<HashSet<ChangeHash>>>>,

    /// `None` for a full state. For a history-only state (loaded for its

    /// changes and tree), the file graphs not built yet, by node: `files`

    /// holds only those loaded with [`State::load_graphs`].

    deferred: Option<crate::fasthash::FastMap<NodeId, Arc<Deferred>>>,

}

use crate::objects::TextRef;

/// A file graph a history-only state has not built: its base from a

/// version 2 checkpoint, if any, and the edits of later changes in the

/// order they were applied (spec §11). Building it later gives the same

/// graph, since graph construction does not depend on order (§8).

///

/// Read from a state index, it stays encoded until first used: a command

/// that builds a few files decodes only their histories.

#[derive(Clone, Default)]

struct Deferred {

    body: std::sync::OnceLock<DeferredBody>,

    /// The body as the index encodes it, decoded into `body` on first use.

    raw: Option<(Arc<[u8]>, std::ops::Range<usize>)>,

    /// Its graph key, once worked out (or read from an index); reset by

    /// any edit added.

    key: std::sync::OnceLock<[u8; 32]>,

}

#[derive(Clone, Default)]

struct DeferredBody {

    base: Option<DeferredBase>,

    edits: Vec<(ChangeHash, Vec<Edit>)>,

    /// The text blob and line count of each insert in `edits` left without

    /// its lines, in order: read when the graph is built.

    texts: Vec<TextRef>,

}

impl Deferred {

    fn from_body(body: DeferredBody) -> Deferred {

        Deferred {

            body: std::sync::OnceLock::from(body),

            ..Deferred::default()

        }

    }

    fn body(&self) -> &DeferredBody {

        self.body.get_or_init(|| match &self.raw {

            // The index was checked against its hash when read, and was

            // written by `write_index`: it decodes.

            Some((bytes, range)) => State::read_deferred_body(&bytes[range.clone()])

                .expect("a state index checked by its hash decodes"),

            None => DeferredBody::default(),

        })

    }

    fn body_mut(&mut self) -> &mut DeferredBody {

        self.body();

        self.raw = None;

        self.body.get_mut().expect("decoded above")

    }

    fn into_body(mut self) -> DeferredBody {

        self.body();

        self.body.take().expect("decoded above")

    }

}

#[derive(Clone)]

enum DeferredBase {

    Inline(Value),

    Blob(ObjectRef),

}

/// A local index of a history-only state (oxygen#41), in the flat layout

/// of [`crate::flat`]: its change metadata, tree (with its view), replace

/// records and every file's deferred history with its graph key, so a

/// command reads it rather than replaying history. Derived only: written

/// from a state loaded the long way, and read back as the same state.

impl State {

    const INDEX_MAGIC: &'static [u8] = b"oxy state index 3\n";

    /// The index of this state, or `None` when it has anything an index

    /// does not hold: a built graph, or text in memory (a change added

    /// whole rather than as its skeleton).

    pub fn write_index(&self) -> Option<Vec<u8>> {

        use crate::flat::Writer;

        let deferred = self.deferred.as_ref()?;

        if !self.files.is_empty() {

            return None;

        }

        let atom = |w: &mut Writer, a: &AtomRef| match a {

            AtomRef::Start => w.u8(0),

            AtomRef::End => w.u8(1),

            AtomRef::Atom { change, index } => {

                w.u8(2);

                w.object(change);

                w.u32(*index);

            }

        };

        // The changes first, as a block a reader can take on its own (see

        // [`State::read_index_changes`]).

        let mut w = Writer::default();

        w.fixed(Self::INDEX_MAGIC);

        let mut changes: Vec<(&ChangeHash, &Arc<ChangeMeta>)> = self.changes.iter().collect();

        changes.sort_by_key(|(h, _)| **h);

        w.count(changes.len());

        for (h, m) in changes {

            w.object(h);

            m.write_flat(&mut w);

        }

        let changes_block = std::mem::take(&mut w);

        w.bytes(&changes_block.0);

        // The tree and the files' histories as blocks of their own, so a

        // reader decodes them side by side.

        let mut tw = Writer::default();

        self.tree.write_flat(&mut tw);

        w.bytes(&tw.0);

        let mut replaces: Vec<(&NodeId, &Vec<ReplaceRecord>)> = self.replaces.iter().collect();

        replaces.sort_by_key(|(n, _)| **n);

        w.count(replaces.len());

        for (n, recs) in replaces {

            w.fixed(&n.0);

            w.count(recs.len());

            // In any order: which record is live does not depend on it.

            let mut recs: Vec<&ReplaceRecord> = recs.iter().collect();

            recs.sort();

            for (c, old, new) in recs {

                w.object(c);

                w.opt_object(old);

                w.opt_object(new);

            }

        }

        let outer = w;

        let mut w = Writer::default();

        let mut nodes: Vec<(&NodeId, &Arc<Deferred>)> = deferred.iter().collect();

        nodes.sort_by_key(|(n, _)| **n);

        w.count(nodes.len());

        for (n, d) in nodes {

            w.fixed(&n.0);

            let mut h = blake3::Hasher::new();

            h.update(b"oxy graph key 1\0");

            h.update(&n.0);

            w.fixed(d.key.get_or_init(|| d.graph_key(h)));

            // Unchanged since read (a change clears `raw`): copied as it is.

            if let Some((bytes, range)) = &d.raw {

                w.bytes(&bytes[range.clone()]);

                continue;

            }

            let d = d.body();

            let outer_w = std::mem::take(&mut w);

            match &d.base {

                None => w.u8(0),

                Some(DeferredBase::Blob(r)) => {

                    w.u8(1);

                    w.object(r);

                }

                Some(DeferredBase::Inline(v)) => {

                    w.u8(2);

                    w.bytes(&cbor::encode(v));

                }

            }

            // Each change's edits with the texts of its inserts, sorted by

            // change: a graph does not depend on the order changes applied

            // in, and the index is then the same however history was read.

            let mut texts = d.texts.iter();

            let mut per_change: Vec<(&ChangeHash, &Vec<Edit>, Vec<&TextRef>)> = Vec::new();

            for (c, edits) in &d.edits {

                let n = edits

                    .iter()

                    .filter(|e| matches!(e, Edit::Insert { lines, .. } if lines.is_empty()))

                    .count();

                per_change.push((c, edits, texts.by_ref().take(n).collect()));

            }

            per_change.sort_by_key(|(c, _, _)| **c);

            w.count(per_change.len());

            for (c, edits, _) in &per_change {

                w.object(c);

                w.count(edits.len());

                for e in edits.iter() {

                    match e {

                        Edit::Insert {

                            node,

                            after,

                            before,

                            lines,

                        } => {

                            if !lines.is_empty() {

                                return None;

                            }

                            w.u8(0);

                            w.fixed(&node.0);

                            atom(&mut w, after);

                            atom(&mut w, before);

                        }

                        Edit::Delete { node, atoms } => {

                            w.u8(1);

                            w.fixed(&node.0);

                            w.count(atoms.len());

                            for r in atoms {

                                w.object(&r.change);

                                w.u32(r.from);

                                w.u32(r.to);

                            }

                        }

                        Edit::Replace { .. } => return None,

                    }

                }

            }

            let texts: Vec<&TextRef> = per_change

                .iter()

                .flat_map(|(_, _, t)| t.iter().copied())

                .collect();

            w.count(texts.len());

            for t in texts {

                match t {

                    TextRef::Blob(r, count, form) => {

                        w.u8(0);

                        w.object(r);

                        w.u32(*count);

                        w.u8(match form {

                            crate::objects::TextForm::Cbor => 0,

                            crate::objects::TextForm::Raw => 1,

                            crate::objects::TextForm::CborFallback => 2,

                        });

                    }

                    TextRef::Inline { change, start, end } => {

                        w.u8(1);

                        w.object(change);

                        w.u32(*start);

                        w.u32(*end);

                    }

                }

            }

            let body = std::mem::replace(&mut w, outer_w);

            w.bytes(&body.0);

        }

        let mut outer = outer;

        outer.bytes(&w.0);

        Some(outer.0)

    }

    /// The state an index holds (see [`State::write_index`]).

    pub fn read_index(bytes: &[u8]) -> Result<State, ObjectError> {

        use crate::flat::Reader;

        let bad = |what: &str| ObjectError::Schema(format!("state index: {what}"));

        let mut r = Reader::new(bytes);

        let mut st = State::new_history_only();

        {

            let _t = crate::timing::phase("index: changes");

            st.changes = Self::read_index_changes(r.bytes()?)?;

        }

        let tree_bytes = r.bytes()?;

        for _ in 0..r.count()? {

            let n = NodeId(r.fixed::<16>()?);

            let mut recs = Vec::new();

            for _ in 0..r.count()? {

                recs.push((r.object()?, r.opt_object()?, r.opt_object()?));

            }

            st.replaces.insert(n, recs);

        }

        let deferred_bytes = r.bytes()?;

        if !r.at_end() {

            return Err(bad("trailing bytes"));

        }

        let (tree, deferred) = rayon::join(

            || TreeGraph::read_flat(&mut Reader::new(tree_bytes)),

            || Self::read_deferred(deferred_bytes),

        );

        st.tree = tree?;

        st.deferred = Some(deferred?);

        Ok(st)

    }

    /// The changes of a state index: its first block, which a reader can

    /// take without the rest (the log needs no more).

    pub fn read_index_changes(

        bytes: &[u8],

    ) -> Result<HashMap<ChangeHash, Arc<ChangeMeta>>, ObjectError> {

        use crate::flat::Reader;

        let mut r = Reader::new(bytes);

        if r.fixed::<18>()?.as_slice() != Self::INDEX_MAGIC {

            return Err(ObjectError::Schema("state index: not a state index".into()));

        }

        let n = r.count()?;

        let mut changes = HashMap::with_capacity(n);

        for _ in 0..n {

            let h = r.object()?;

            changes.insert(h, Arc::new(ChangeMeta::read_flat(&mut r)?));

        }

        if !r.at_end() {

            return Err(ObjectError::Schema("state index: trailing bytes".into()));

        }

        Ok(changes)

    }

    /// The deferred block: each node with its graph key and its body,

    /// left encoded (see [`Deferred`]).

    fn read_deferred(

        bytes: &[u8],

    ) -> Result<crate::fasthash::FastMap<NodeId, Arc<Deferred>>, ObjectError> {

        use crate::flat::Reader;

        let _t = crate::timing::phase("index: deferred");

        let all: Arc<[u8]> = Arc::from(bytes);

        let mut r = Reader::new(bytes);

        let mut deferred = crate::fasthash::FastMap::default();

        let n = r.count()?;

        deferred.reserve(n);

        for _ in 0..n {

            let id = NodeId(r.fixed::<16>()?);

            let key = r.fixed::<32>()?;

            let body = r.bytes()?;

            let start = body.as_ptr() as usize - bytes.as_ptr() as usize;

            let d = Deferred {

                raw: Some((all.clone(), start..start + body.len())),

                key: std::sync::OnceLock::from(key),

                ..Deferred::default()

            };

            deferred.insert(id, Arc::new(d));

        }

        if !r.at_end() {

            return Err(ObjectError::Schema("state index: trailing bytes".into()));

        }

        Ok(deferred)

    }

    /// One node's deferred history as the index encodes it.

    fn read_deferred_body(bytes: &[u8]) -> Result<DeferredBody, ObjectError> {

        use crate::flat::Reader;

        let bad = |what: &str| ObjectError::Schema(format!("state index: {what}"));

        let node = |r: &mut Reader| -> Result<NodeId, ObjectError> { Ok(NodeId(r.fixed::<16>()?)) };

        let atom = |r: &mut Reader| -> Result<AtomRef, ObjectError> {

            Ok(match r.u8()? {

                0 => AtomRef::Start,

                1 => AtomRef::End,

                2 => AtomRef::Atom {

                    change: r.object()?,

                    index: r.u32()?,

                },

                _ => return Err(bad("atom")),

            })

        };

        let mut r = Reader::new(bytes);

        let base = match r.u8()? {

            0 => None,

            1 => Some(DeferredBase::Blob(r.object()?)),

            _ => Some(DeferredBase::Inline(cbor::decode(r.bytes()?)?)),