//! Rewrite propagation (spec §8.4): when change `old` is replaced by `new`,
//! every change that depends on `old` is rewritten with its references
//! translated through a correspondence from `old`'s atoms and tree operations
//! to `new`'s. Rewritten changes keep their ChangeIds. Where a reference has
//! no counterpart, the rewrite fails with the orphaned reference named;
//! turning that into a reified conflict is future work.

use crate::hash::ChangeHash;
use crate::objects::{AtomRange, AtomRef, Change, Edit, OpRef, TreeOp};
use crate::state::State;
use std::collections::{BTreeSet, HashMap};

/// How `old`'s atoms and tree operations map onto `new`'s.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Correspondence {
    /// Old atom index → new atom index, per node the old change edited.
    /// Absent means the atom no longer exists.
    pub atoms: HashMap<u32, u32>,
    /// Old tree-op index → new tree-op index.
    pub ops: HashMap<u32, u32>,
}

impl Correspondence {
    /// Every atom and operation keeps its index: `new` has the same edits as `old`.
    pub fn identity(change: &Change) -> Correspondence {
        let mut atoms = HashMap::new();
        let mut per_node: HashMap<crate::ids::NodeId, u32> = HashMap::new();
        for e in &change.edits {
            if let Edit::Insert { node, lines, .. } = e {
                let idx = per_node.entry(*node).or_default();
                for _ in lines {
                    atoms.insert(*idx, *idx);
                    *idx += 1;
                }
            }
        }
        // Atom indices are per node, so the identity map is index → index for
        // the largest index any node uses.
        let max = per_node.values().copied().max().unwrap_or(0);
        for i in 0..max {
            atoms.insert(i, i);
        }
        let ops = (0..change.tree_ops.len() as u32).map(|i| (i, i)).collect();
        Correspondence { atoms, ops }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum RewriteError {
    #[error("change {change} references {what} of the rewritten change, which no longer exists")]
    Orphaned { change: ChangeHash, what: String },
    #[error("change {0} is not in the state")]
    Missing(ChangeHash),
}

/// What a rewritten change was replaced by.
#[derive(Clone, Debug)]
enum Replacement {
    Change(ChangeHash),
    /// Removed: dependents inherit these dependencies instead.
    Removed(Vec<ChangeHash>),
}

/// Rewrite every change in `state` that depends on `old`, given `new` (or
/// `None` when `old` is removed) and the correspondence between them.
/// Returns `(old hash, rewritten change)` pairs in dependency order, not
/// including `old` → `new` itself.
///
/// `load` fetches a change's full object (a state keeps only metadata).
pub fn propagate(state: &State, old: ChangeHash, new: Option<&Change>, corr: &Correspondence, load: &dyn Fn(&ChangeHash) -> Option<Change>) -> Result<Vec<(ChangeHash, Change)>, RewriteError> {
    let old_change = state.get(&old).ok_or(RewriteError::Missing(old))?;
    let mut mapping: HashMap<ChangeHash, Replacement> = HashMap::new();
    mapping.insert(old, match new {
        Some(n) => Replacement::Change(n.hash()),
        None => Replacement::Removed(old_change.deps.clone()),
    });
    let mut corrs: HashMap<ChangeHash, Correspondence> = HashMap::new();
    corrs.insert(old, corr.clone());

    // Dependents in topological order: every change whose closure has `old`,
    // sorted so that dependencies come first.
    let mut dependents: Vec<ChangeHash> = state.changes().filter(|(h, _)| **h != old && state.closure(h).contains(&old)).map(|(h, _)| *h).collect();
    dependents.sort_by_key(|h| (state.closure(h).len(), *h));

    let mut out = Vec::new();
    for h in dependents {
        let c = load(&h).ok_or(RewriteError::Missing(h))?;
        let rewritten = rewrite_one(&h, &c, &mapping, &corrs)?;
        let rh = rewritten.hash();
        corrs.insert(h, Correspondence::identity(&rewritten));
        mapping.insert(h, Replacement::Change(rh));
        out.push((h, rewritten));
    }
    Ok(out)
}

fn rewrite_one(hash: &ChangeHash, c: &Change, mapping: &HashMap<ChangeHash, Replacement>, corrs: &HashMap<ChangeHash, Correspondence>) -> Result<Change, RewriteError> {
    let orphan = |what: String| RewriteError::Orphaned { change: *hash, what };
    let target = |change: &ChangeHash, what: &str| -> Result<Option<ChangeHash>, RewriteError> {
        match mapping.get(change) {
            Some(Replacement::Change(nh)) => Ok(Some(*nh)),
            Some(Replacement::Removed(_)) => Err(orphan(what.to_string())),
            None => Ok(None),
        }
    };
    let map_atom = |r: AtomRef| -> Result<AtomRef, RewriteError> {
        match r {
            AtomRef::Atom { change, index } => match target(&change, &format!("atom {index}"))? {
                Some(nh) => {
                    let idx = corrs[&change].atoms.get(&index).ok_or_else(|| orphan(format!("atom {index}")))?;
                    Ok(AtomRef::Atom { change: nh, index: *idx })
                }
                None => Ok(r),
            },
            other => Ok(other),
        }
    };
    let map_op = |r: OpRef| -> Result<OpRef, RewriteError> {
        match target(&r.change, &format!("tree operation {}", r.index))? {
            Some(nh) => {
                let idx = corrs[&r.change].ops.get(&r.index).ok_or_else(|| orphan(format!("tree operation {}", r.index)))?;
                Ok(OpRef { change: nh, index: *idx })
            }
            None => Ok(r),
        }
    };
    let mut out = c.clone();
    let mut deps: BTreeSet<ChangeHash> = BTreeSet::new();
    for d in &c.deps {
        match mapping.get(d) {
            Some(Replacement::Change(nh)) => {
                deps.insert(*nh);
            }
            Some(Replacement::Removed(inherited)) => deps.extend(inherited.iter().copied()),
            None => {
                deps.insert(*d);
            }
        }
    }
    out.deps = deps.into_iter().collect();
    for op in &mut out.tree_ops {
        match op {
            TreeOp::RemoveNode { placement, .. } => *placement = map_op(*placement)?,
            TreeOp::Move { old, .. } => *old = map_op(*old)?,
            TreeOp::SetMode { old, .. } => *old = map_op(*old)?,
            TreeOp::AddNode { .. } => {}
        }
    }
    for e in &mut out.edits {
        match e {
            Edit::Insert { after, before, .. } => {
                *after = map_atom(*after)?;
                *before = map_atom(*before)?;
            }
            Edit::Delete { atoms, .. } => {
                let mut mapped: Vec<(ChangeHash, u32)> = Vec::new();
                for r in atoms.iter() {
                    for i in r.from..r.to {
                        match map_atom(AtomRef::Atom { change: r.change, index: i })? {
                            AtomRef::Atom { change, index } => mapped.push((change, index)),
                            _ => unreachable!(),
                        }
                    }
                }
                let mut ranges: Vec<AtomRange> = Vec::new();
                for (ch, i) in mapped {
                    match ranges.last_mut() {
                        Some(r) if r.change == ch && r.to == i => r.to = i + 1,
                        _ => ranges.push(AtomRange { change: ch, from: i, to: i + 1 }),
                    }
                }
                *atoms = ranges;
            }
            Edit::Replace { .. } => {}
        }
    }
    out.normalize();
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hash::ObjectRef;
    use crate::ids::{ChangeId, NodeId};
    use crate::objects::{ChangeKind, Hlc, Mode, NodeKind};
    use crate::state::FileContent;

    fn change(deps: Vec<ChangeHash>, tree_ops: Vec<TreeOp>, edits: Vec<Edit>, msg: &str) -> Change {
        let mut c = Change { change_id: ChangeId::random(), deps, author: ObjectRef::ZERO, hlc: Hlc { wall_ms: 1, logical: 0 }, message: msg.into(), tree_ops, edits, kind: ChangeKind::Normal };
        c.normalize();
        c
    }

    #[test]
    fn describe_propagates_through_dependents() {
        let mut st = State::new();
        let f = NodeId::random();
        let c1 = change(vec![], vec![TreeOp::AddNode { node: f, kind: NodeKind::File, mode: Mode::TEXT, parent: NodeId::ROOT, name: "f".into() }], vec![Edit::Insert { node: f, after: AtomRef::Start, before: AtomRef::End, lines: vec![b"a\n".to_vec(), b"b\n".to_vec()] }], "one");
        let h1 = st.add(c1.clone()).unwrap();
        let c2 = change(vec![h1], vec![], vec![Edit::Delete { node: f, atoms: vec![AtomRange { change: h1, from: 0, to: 1 }] }, Edit::Insert { node: f, after: AtomRef::Atom { change: h1, index: 1 }, before: AtomRef::End, lines: vec![b"c\n".to_vec()] }], "two");
        let h2 = st.add(c2.clone()).unwrap();
        let c3 = change(vec![h2], vec![TreeOp::SetMode { node: f, old: OpRef { change: h1, index: 0 }, mode: Mode { exec: true, opaque: false } }], vec![Edit::Insert { node: f, after: AtomRef::Atom { change: h2, index: 0 }, before: AtomRef::End, lines: vec![b"d\n".to_vec()] }], "three");
        let h3 = st.add(c3.clone()).unwrap();
        let before = st.materialize();

        // Rewrite c1's message; everything downstream follows.
        let mut n1 = c1.clone();
        n1.message = "one, described".into();
        let all: HashMap<ChangeHash, Change> = [(h1, c1.clone()), (h2, c2.clone()), (h3, c3.clone())].into_iter().collect();
        let load = |h: &ChangeHash| all.get(h).cloned();
        let rewritten = propagate(&st, h1, Some(&n1), &Correspondence::identity(&c1), &load).unwrap();
        assert_eq!(rewritten.iter().map(|(h, _)| *h).collect::<Vec<_>>(), vec![h2, h3]);
        let mut st2 = State::new();
        st2.add(n1.clone()).unwrap();
        for (_, c) in &rewritten {
            st2.add(c.clone()).unwrap();
        }
        assert_eq!(st2.len(), 3);
        let after = st2.materialize();
        assert_eq!(after.files["f"].mode, before.files["f"].mode);
        match (&after.files["f"].content, &before.files["f"].content) {
            (FileContent::Text { bytes: a, .. }, FileContent::Text { bytes: b, .. }) => assert_eq!(a, b),
            _ => panic!(),
        }
        assert!(matches!(&after.files["f"].content, FileContent::Text { bytes, .. } if bytes == b"b\nc\nd\n"));
        // ChangeIds survive; hashes change; deps point at the new hashes.
        assert_eq!(rewritten[0].1.change_id, c2.change_id);
        assert_eq!(rewritten[0].1.deps, vec![n1.hash()]);
        assert_ne!(rewritten[0].1.hash(), h2);
        assert_eq!(rewritten[1].1.deps, vec![rewritten[0].1.hash()]);
    }

    #[test]
    fn missing_counterpart_is_reported() {
        let mut st = State::new();
        let f = NodeId::random();
        let c1 = change(vec![], vec![TreeOp::AddNode { node: f, kind: NodeKind::File, mode: Mode::TEXT, parent: NodeId::ROOT, name: "f".into() }], vec![Edit::Insert { node: f, after: AtomRef::Start, before: AtomRef::End, lines: vec![b"a\n".to_vec(), b"b\n".to_vec()] }], "one");
        let h1 = st.add(c1.clone()).unwrap();
        let c2 = change(vec![h1], vec![], vec![Edit::Insert { node: f, after: AtomRef::Atom { change: h1, index: 1 }, before: AtomRef::End, lines: vec![b"c\n".to_vec()] }], "two");
        let h2 = st.add(c2.clone()).unwrap();
        let all: HashMap<ChangeHash, Change> = [(h1, c1.clone()), (h2, c2)].into_iter().collect();
        let load = |h: &ChangeHash| all.get(h).cloned();
        // The rewritten c1 drops line "b": atom 1 has no counterpart.
        let mut n1 = c1.clone();
        n1.edits = vec![Edit::Insert { node: f, after: AtomRef::Start, before: AtomRef::End, lines: vec![b"a\n".to_vec()] }];
        let corr = Correspondence { atoms: [(0, 0)].into_iter().collect(), ops: [(0, 0)].into_iter().collect() };
        let err = propagate(&st, h1, Some(&n1), &corr, &load).unwrap_err();
        assert_eq!(err, RewriteError::Orphaned { change: h2, what: "atom 1".into() });
        assert_eq!(propagate(&st, ObjectRef([9; 32]), Some(&n1), &corr, &load).unwrap_err(), RewriteError::Missing(ObjectRef([9; 32])));
    }

    #[test]
    fn removal_rewires_dependencies_or_fails_on_references() {
        let mut st = State::new();
        let f = NodeId::random();
        let g = NodeId::random();
        let c1 = change(vec![], vec![TreeOp::AddNode { node: f, kind: NodeKind::File, mode: Mode::TEXT, parent: NodeId::ROOT, name: "f".into() }], vec![], "one");
        let h1 = st.add(c1.clone()).unwrap();
        // c2 merely depends on c1 explicitly (touches another file).
        let c2 = change(vec![h1], vec![TreeOp::AddNode { node: g, kind: NodeKind::File, mode: Mode::TEXT, parent: NodeId::ROOT, name: "g".into() }], vec![], "two");
        let h2 = st.add(c2.clone()).unwrap();
        let c3 = change(vec![h2], vec![], vec![Edit::Insert { node: g, after: AtomRef::Start, before: AtomRef::End, lines: vec![b"x\n".to_vec()] }], "three");
        let h3 = st.add(c3.clone()).unwrap();
        let mut all: HashMap<ChangeHash, Change> = [(h1, c1), (h2, c2), (h3, c3)].into_iter().collect();
        let load_all = all.clone();
        let load = move |h: &ChangeHash| load_all.get(h).cloned();
        // Removing c2 (an explicit-only dependency of nothing that references it) works: c3 inherits h1.
        let out = propagate(&st, h2, None, &Correspondence::default(), &load).unwrap();
        assert_eq!(out.len(), 1);
        assert_eq!(out[0].0, h3);
        assert_eq!(out[0].1.deps, vec![h1]);
        // Removing c1 fails: c2 does not reference c1 (only an explicit dep), but
        // nothing else does either, so it succeeds and c2 inherits no deps.
        let out = propagate(&st, h1, None, &Correspondence::default(), &load).unwrap();
        assert!(out.iter().any(|(h, c)| *h == h2 && c.deps.is_empty()));
        // Removing c2 when c3 references its placement fails.
        let c4 = change(vec![h2], vec![TreeOp::RemoveNode { node: g, placement: OpRef { change: h2, index: 0 } }], vec![], "four");
        let h4 = st.add(c4.clone()).unwrap();
        all.insert(h4, c4);
        let load2 = |h: &ChangeHash| all.get(h).cloned();
        let err = propagate(&st, h2, None, &Correspondence::default(), &load2).unwrap_err();
        assert!(matches!(err, RewriteError::Orphaned { change, .. } if change == h4 || change == h3), "{err:?}");
    }
}