//! 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 {

    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: HashMap<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<HashMap<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).

#[derive(Clone, Default)]

struct Deferred {

    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>,

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

    /// any edit added.

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

}

#[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 1\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);

            }

        };

        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);

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

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

            for d in &m.deps {

                w.object(d);

            }

            w.object(&m.author);

            w.u64(m.hlc.wall_ms);

            w.u32(m.hlc.logical);

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

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

        }

        // 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);

            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 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)));

        }

        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);

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

            return Err(bad("not a state index"));

        }

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

        let mut st = State::new_history_only();

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

        let n = r.count()?;

        st.changes.reserve(n);

        for _ in 0..n {

            let h = r.object()?;

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

            let mut deps = Vec::new();

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

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

            }

            let meta = 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)?,

            };

            st.changes.insert(h, Arc::new(meta));

        }

        drop(t);

        let tree_bytes = r.bytes()?;

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

            let n = node(&mut r)?;

            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)

    }

    fn read_deferred(bytes: &[u8]) -> Result<HashMap<NodeId, Arc<Deferred>>, ObjectError> {

        use crate::flat::Reader;

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

        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 mut deferred: HashMap<NodeId, Arc<Deferred>> = HashMap::new();

        let n = r.count()?;

        deferred.reserve(n);

        for _ in 0..n {

            let id = node(&mut r)?;

            let base = match r.u8()? {

                0 => None,

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

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

            };

            let mut edits = Vec::new();

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

                let c = r.object()?;

                let mut es = Vec::new();

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

                    es.push(match r.u8()? {

                        0 => Edit::Insert {

                            node: node(&mut r)?,

                            after: atom(&mut r)?,

                            before: atom(&mut r)?,

                            lines: crate::lines::Lines::default(),

                        },

                        1 => {

                            let node = node(&mut r)?;

                            let mut atoms = Vec::new();

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

                                atoms.push(crate::objects::AtomRange {

                                    change: r.object()?,

                                    from: r.u32()?,

                                    to: r.u32()?,

                                });

                            }

                            Edit::Delete { node, atoms }

                        }

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

                    });

                }

                edits.push((c, es));

            }

            let mut texts = Vec::new();

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

                texts.push(match r.u8()? {

                    0 => TextRef::Blob(

                        r.object()?,

                        r.u32()?,

                        match r.u8()? {

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

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

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

                        },

                    ),

                    _ => TextRef::Inline {

                        change: r.object()?,

                        start: r.u32()?,

                        end: r.u32()?,

                    },

                });

            }

            let d = Deferred {

                base,

                edits,

                texts,

                key: std::sync::OnceLock::from(r.fixed::<32>()?),

            };

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

        }

        if !r.at_end() {

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

        }

        Ok(deferred)

    }

}

impl Clone for State {

    fn clone(&self) -> State {

        State {

            changes: self.changes.clone(),

            tree: self.tree.clone(),

            files: self.files.clone(),

            replaces: self.replaces.clone(),

            // A cache: the clone builds its own.

            closure: Mutex::new(HashMap::new()),

            deferred: self.deferred.clone(),

        }

    }

}

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

pub enum FileContent {

    Text {

        bytes: Vec<u8>,

        materialized: Materialized,

    },

    /// Live blobs; more than one is a replace conflict (spec §9.3).

    Opaque {

        blobs: Vec<ObjectRef>,

    },

    Dir,

}

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

pub struct FileState {

    pub node: NodeId,

    pub kind: NodeKind,

    pub mode: Mode,

    pub content: FileContent,

}

impl FileState {

    pub fn is_conflicted(&self) -> bool {

        match &self.content {

            FileContent::Text { materialized, .. } => materialized.is_conflicted(),