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

    changes: HashMap<ChangeHash, 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, Deferred>>,

}

/// A stored insert's text blob and line count.

type TextRef = (ObjectRef, u32);

/// 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<(ObjectRef, u32)>,

}

#[derive(Clone)]

enum DeferredBase {

    Inline(Value),

    Blob(ObjectRef),

}

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

            closure: Mutex::new(self.closure.lock().unwrap().clone()),

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

            FileContent::Opaque { blobs } => blobs.len() > 1,

            FileContent::Dir => false,

        }

    }

}

/// See `State::conflict_regions`.

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

pub struct ConflictRegions {

    pub plain: Vec<(AtomKey, Vec<u8>)>,

    pub region_atoms: Vec<AtomKey>,

    pub contributors: Vec<ChangeHash>,

}

/// The whole tree of a state, materialized.

#[derive(Clone, Debug, Default)]

pub struct Snapshot {

    pub tree: TreeView,

    pub files: BTreeMap<String, FileState>,

    /// Text files left unrendered by [`State::materialize_except`]: their

    /// `bytes` are empty and they are known not to be conflicted.

    pub unrendered: BTreeSet<String>,

}

impl Snapshot {

    pub fn is_conflicted(&self) -> bool {

        self.tree.is_conflicted() || self.files.values().any(|f| f.is_conflicted())

    }

}

fn vertex(r: AtomRef) -> Vertex {

    match r {

        AtomRef::Start => Vertex::Start,

        AtomRef::End => Vertex::End,

        AtomRef::Atom { change, index } => Vertex::Atom((change, index)),

    }

}

/// Replace self references in a tree op with the concrete change hash.

pub fn resolve_tree_op(op: &TreeOp, me: ChangeHash) -> TreeOp {

    match op {

        TreeOp::RemoveNode { node, placement } => TreeOp::RemoveNode {

            node: *node,

            placement: placement.resolve(me),

        },

        TreeOp::Move {

            node,

            old,

            parent,

            name,

        } => TreeOp::Move {

            node: *node,

            old: old.resolve(me),

            parent: *parent,

            name: name.clone(),

        },

        TreeOp::SetMode { node, old, mode } => TreeOp::SetMode {

            node: *node,

            old: old.resolve(me),

            mode: *mode,

        },

        add @ TreeOp::AddNode { .. } => add.clone(),

    }

}

/// Apply one change's edits of one file (inserts and deletes; replaces

/// live beside the graphs) to that file's graph.

fn apply_file_edits(g: &mut FileGraph, h: ChangeHash, edits: Vec<Edit>) {

    let mut next = 0u32;

    for e in edits {

        match e {

            Edit::Insert {

                after,

                before,

                lines,

                ..

            } => {

                let first = next;

                next += lines.len() as u32;

                g.insert_owned(

                    h,

                    first,

                    vertex(after.resolve(h)),

                    vertex(before.resolve(h)),

                    lines,

                );

            }

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

                for r in atoms {

                    g.delete(h, &r.resolve(h));

                }

            }

            Edit::Replace { .. } => {}

        }

    }

}

impl State {

    pub fn new() -> Self {

        Self::default()

    }

    /// An empty state for history only: changes and tree are kept, file

    /// edits are recorded but no file graph is built until

    /// [`State::load_graphs`] asks for it. For log and other commands that

    /// need no file contents, or only a few files.

    pub fn new_history_only() -> Self {

        State {

            deferred: Some(HashMap::new()),

            ..Self::default()

        }

    }

    /// Whether this state builds file graphs only when asked

    /// ([`State::new_history_only`]), even if it has none waiting.

    pub fn is_history_only(&self) -> bool {

        self.deferred.is_some()

    }

    /// Whether this state holds every file graph.

    pub fn is_full(&self) -> bool {

        self.deferred.as_ref().is_none_or(|d| d.is_empty())

    }

    /// The graph of `node`, which must be loaded: asking a history-only

    /// state for a graph it has not built is a bug, never an empty file.

    fn graph(&self, node: &NodeId) -> Option<&FileGraph> {

        if let Some(d) = &self.deferred {

            assert!(

                !d.contains_key(node),

                "file graph {} of a history-only state was not loaded (State::load_graphs)",

                node.short(12)

            );

        }

        self.files.get(node).map(|g| g.as_ref())

    }

    /// Operations over every file need every graph.

    fn require_full(&self, what: &str) {

        assert!(

            self.is_full(),

            "{what} needs every file graph; this state was loaded for history only"

        );

    }

    /// Build the graphs of `nodes` in a history-only state: from their

    /// checkpoint base (blobs read in one batch through `load`) and their

    /// deferred edits. Nodes already built, or with no graph, are skipped.

    pub fn load_graphs(

        &mut self,

        nodes: &[NodeId],

        load: impl FnOnce(&[ObjectRef]) -> Result<Vec<Vec<u8>>, ObjectError>,

    ) -> Result<(), ObjectError> {

        let Some(deferred) = self.deferred.as_mut() else {

            return Ok(());

        };

        let mut taken: Vec<(NodeId, Deferred)> = nodes

            .iter()

            .filter_map(|n| deferred.remove(n).map(|d| (*n, d)))

            .collect();

        let mut refs: Vec<ObjectRef> = taken

            .iter()

            .filter_map(|(_, d)| match &d.base {

                Some(DeferredBase::Blob(r)) => Some(*r),

                _ => None,

            })

            .collect();

        // Text blobs of inserts read without their lines, after the bases.

        let bases = refs.len();

        refs.extend(

            taken

                .iter()

                .flat_map(|(_, d)| d.texts.iter().map(|(r, _)| *r)),

        );

        let mut blobs = if refs.is_empty() {

            Vec::new()

        } else {

            load(&refs)?

        };

        if blobs.len() != refs.len() {

            return Err(ObjectError::Schema("graph blobs: wrong count".into()));

        }

        let mut texts = blobs.split_off(bases).into_iter();

        let mut blobs = blobs.into_iter();

        // Pair each graph with its blob, then build them all in parallel:

        // graphs are independent.

        let mut work: Vec<(NodeId, Deferred, Option<Vec<u8>>)> = Vec::with_capacity(taken.len());

        for (node, mut d) in taken.drain(..) {

            // Fill in the lines of inserts read without them, in order.

            let mut want = std::mem::take(&mut d.texts).into_iter();

            for (_, edits) in d.edits.iter_mut() {

                for e in edits.iter_mut() {

                    if let Edit::Insert { lines, .. } = e {

                        if lines.is_empty() {

                            let (_, count) = want.next().ok_or_else(|| {

                                ObjectError::Schema("an insert's text is missing".into())

                            })?;

                            let content = texts.next().ok_or_else(|| {

                                ObjectError::Schema("text blobs: wrong count".into())

                            })?;

                            *lines = crate::objects::text_lines(&content, count)?;

                        }

                    }

                }

            }

            let bytes = match &d.base {

                Some(DeferredBase::Blob(_)) => Some(

                    blobs

                        .next()

                        .ok_or_else(|| ObjectError::Schema("graph blobs: wrong count".into()))?,

                ),

                _ => None,

            };

            work.push((node, d, bytes));

        }

        use rayon::prelude::*;

        let built: Vec<(NodeId, FileGraph)> = work

            .into_par_iter()

            .map(|(node, mut d, bytes)| {

                let mut graph = match (d.base.take(), bytes) {

                    (Some(DeferredBase::Inline(v)), _) => FileGraph::from_value(&v)?,

                    (Some(DeferredBase::Blob(_)), Some(bytes)) => {

                        FileGraph::from_value(&cbor::decode(&bytes)?)?

                    }

                    _ => FileGraph::default(),

                };

                for (h, edits) in std::mem::take(&mut d.edits) {

                    apply_file_edits(&mut graph, h, edits);

                }

                Ok((node, graph))

            })

            .collect::<Result<_, ObjectError>>()?;

        self.files

            .extend(built.into_iter().map(|(n, g)| (n, Arc::new(g))));

        Ok(())

    }

    /// Drop the built graphs of `nodes`, for a caller rendering a large

    /// tree a batch at a time. The state can no longer render those files,

    /// so it must not be used for them again.

    pub fn drop_graphs(&mut self, nodes: &[NodeId]) {

        for n in nodes {

            self.files.remove(n);

        }

    }

    /// Build every graph a history-only state has not built, making it full.

    pub fn load_all_graphs(

        &mut self,

        load: impl FnOnce(&[ObjectRef]) -> Result<Vec<Vec<u8>>, ObjectError>,

    ) -> Result<(), ObjectError> {

        let nodes: Vec<NodeId> = match &self.deferred {

            Some(d) => d.keys().copied().collect(),

            None => return Ok(()),

        };

        self.load_graphs(&nodes, load)?;

        self.deferred = None;

        Ok(())

    }

    /// A key for the graph of `node` in a history-only state that has not

    /// built it: equal keys mean equal graphs, so equal renderings (a

    /// rendering depends only on the graph and on the changes' own

    /// immutable metadata). It names the node, its checkpoint base (a blob,

    /// an inline value, or none) and the set of later changes that edited

    /// the file.

    /// `None` for a full state or a graph already built.

    pub fn graph_key(&self, node: &NodeId) -> Option<[u8; 32]> {

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

        if self.files.contains_key(node) {

            return None;

        }

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

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

        // A change's edits are applied per node: the node is part of it.

        h.update(&node.0);

        let Some(d) = deferred.get(node) else {

            // Never edited: the empty graph.

            h.update(&[0]);

            return Some(*h.finalize().as_bytes());

        };

        match &d.base {

            None => h.update(&[0]),

            Some(DeferredBase::Blob(r)) => h.update(&[1]).update(&r.0),

            Some(DeferredBase::Inline(v)) => h.update(&[2]).update(&cbor::encode(v)),

        };

        let mut changes: Vec<&ChangeHash> = d.edits.iter().map(|(c, _)| c).collect();

        changes.sort();

        changes.dedup();

        h.update(&(changes.len() as u64).to_le_bytes());

        for c in changes {

            h.update(&c.0);

        }

        Some(*h.finalize().as_bytes())

    }

    pub fn len(&self) -> usize {

        self.changes.len()

    }

    pub fn is_empty(&self) -> bool {

        self.changes.is_empty()

    }

    pub fn contains(&self, h: &ChangeHash) -> bool {

        self.changes.contains_key(h)

    }

    pub fn get(&self, h: &ChangeHash) -> Option<&ChangeMeta> {

        self.changes.get(h)

    }

    pub fn changes(&self) -> impl Iterator<Item = (&ChangeHash, &ChangeMeta)> {

        self.changes.iter()

    }

    pub fn tree(&self) -> &TreeGraph {

        &self.tree

    }

    pub fn file(&self, node: &NodeId) -> Option<&FileGraph> {

        self.graph(node)

    }

    /// Add a change whose dependencies are already present.

    pub fn add(&mut self, change: Change) -> Result<ChangeHash, StateError> {

        let h = change.hash();

        if self.changes.contains_key(&h) {

            return Ok(h);

        }

        for d in &change.deps {

            if !self.changes.contains_key(d) {

                return Err(StateError::MissingDep(*d, h));

            }

        }

        Ok(self.add_unchecked(change))

    }

    /// Add a change in any order; call `validate` once all are in.

    pub fn add_unchecked(&mut self, change: Change) -> ChangeHash {

        let h = change.hash();

        self.add_known(h, change, Vec::new());

        h

    }

    /// Add a change read from the store under hash `h`, without hashing it

    /// again (for a version 2 change that means chunking all its text).

    /// `texts` names, in order, the inserts left without their lines (see

    /// [`crate::store::Store::get_change_lazy`]); only a history-only state

    /// takes them, and only for graphs it has not built.

    pub fn add_known(&mut self, h: ChangeHash, change: Change, texts: Vec<(ObjectRef, u32)>) {

        if self.changes.contains_key(&h) {

            return;

        }

        let meta = ChangeMeta::of(&change);

        self.apply_with(h, change, texts);

        self.changes.insert(h, meta);

        self.closure.lock().unwrap().clear();

    }

    // ------------------------------------------------------------ checkpoints (spec §11)