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

}

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

}

#[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 { .. } => {}

        }

    }

}

/// Build the graphs of deferred entries: their checkpoint bases and the

/// texts of inserts read without their lines are loaded in one call, then

/// the graphs are built in parallel (they are independent).

fn build_deferred(

    mut taken: Vec<(NodeId, Deferred)>,

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

    read_change: impl Fn(&ChangeHash) -> Result<Arc<Vec<u8>>, ObjectError>,

) -> Result<Vec<(NodeId, FileGraph)>, ObjectError> {

    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().filter_map(|t| match t {

            TextRef::Blob(r, _) => Some(*r),

            TextRef::Inline { .. } => None,

        })

    }));

    // And the encodings of changes whose inline inserts are needed,

    // each read once.

    let mut changes: HashMap<ChangeHash, Arc<Vec<u8>>> = HashMap::new();

    for t in taken.iter().flat_map(|(_, d)| d.texts.iter()) {

        if let TextRef::Inline { change, .. } = t {

            if !changes.contains_key(change) {

                changes.insert(*change, read_change(change)?);

            }

        }

    }

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

                        *lines = match want.next() {

                            Some(TextRef::Blob(_, count)) => {

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

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

                                })?;

                                crate::objects::text_lines(&content, count)?

                            }

                            Some(TextRef::Inline { change, start, end }) => {

                                let bytes = &changes[&change];

                                let range =

                                    bytes.get(start as usize..end as usize).ok_or_else(|| {

                                        ObjectError::Schema(

                                            "an inline insert lies outside its change".into(),

                                        )

                                    })?;

                                crate::objects::inline_lines(range)?

                            }

                            None => {

                                return Err(ObjectError::Schema(

                                    "an insert's text is missing".into(),

                                ))

                            }

                        };

                    }

                }

            }

        }

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

    Ok(built)

}

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

        read_change: impl Fn(&ChangeHash) -> Result<Arc<Vec<u8>>, ObjectError>,

    ) -> Result<(), ObjectError> {

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

            return Ok(());

        };

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

            .iter()

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

            .collect();

        let built = build_deferred(taken, load, read_change)?;

        self.files

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

        Ok(())

    }

    /// The graphs of `nodes`, built from this history-only state's deferred

    /// edits without building them in it: the state stays history only, so

    /// changes read without their lines can still be added to it and every

    /// graph keeps its key. For a caller that renders a few files of many

    /// states in turn (verifying an import).

    pub fn build_graphs_detached(

        &self,

        nodes: &[NodeId],

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

        read_change: impl Fn(&ChangeHash) -> Result<Arc<Vec<u8>>, ObjectError>,

    ) -> Result<Vec<(NodeId, FileGraph)>, ObjectError> {

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

            return Ok(Vec::new());

        };

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

            .iter()

            .map(|n| (*n, deferred.get(n).cloned().unwrap_or_default()))

            .collect();

        build_deferred(copies, load, read_change)

    }

    /// Render a graph built for this state (see

    /// [`State::build_graphs_detached`]), as [`State::render_file`] does.

    pub fn render_graph(&self, g: &FileGraph) -> (Vec<u8>, Materialized) {

        let m = g.materialize(self);

        let bytes = g.render(&m, &|h| self.label(h));

        (bytes, m)

    }

    /// Every file graph a history-only state has not built, deleted files'

    /// included.

    pub fn unbuilt_nodes(&self) -> Vec<NodeId> {

        self.deferred

            .as_ref()

            .map(|d| d.keys().copied().collect())

            .unwrap_or_default()

    }

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

        read_change: impl Fn(&ChangeHash) -> Result<Arc<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, read_change)?;

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