//! Three-way merges: text line by line, JSON by key, and whole trees.
//!
//! The orientation is the same everywhere: `ours` is main and `theirs` is
//! the change being landed or synced. On the server ([`Mode::Land`]) any
//! conflict stops the landing. In a workspace ([`Mode::Sync`]) text conflicts
//! are written into the file with markers, and conflicts that can't be marked
//! (binary files, a file deleted on one side) keep the change's side and say
//! so.

use std::collections::BTreeMap;

use serde_json::Value;

use crate::chunk;
use crate::diff::{Hunk, hunks, is_binary, lines};
use crate::object::{Kind, Object};
use crate::tree::{self, EntryKind, TreeEntry};
use crate::{Id, Layered, MemSource, Result, Source};

/// The result of a text merge.
pub struct TextMerge {
    pub content: Vec<u8>,
    pub conflicts: usize,
}

/// Line-by-line three-way merge. Changes on both sides to the same or
/// touching lines are a conflict unless they're identical.
pub fn text(base: &[u8], ours: &[u8], theirs: &[u8], ours_label: &str, theirs_label: &str) -> TextMerge {
    let b = lines(base);
    let o = lines(ours);
    let t = lines(theirs);
    let ho = hunks(base, ours);
    let ht = hunks(base, theirs);
    let mut out = Vec::with_capacity(ours.len().max(theirs.len()));
    let mut conflicts = 0;
    let (mut i, mut j) = (0, 0);
    let mut pos: u32 = 0;
    while i < ho.len() || j < ht.len() {
        // Start a group at the earliest hunk.
        let take_ours = j >= ht.len() || (i < ho.len() && ho[i].before.start <= ht[j].before.start);
        let first = if take_ours { &ho[i] } else { &ht[j] };
        let lo = first.before.start;
        let mut hi = first.before.end;
        let (gi, gj) = (i, j);
        if take_ours {
            i += 1;
        } else {
            j += 1;
        }
        // Absorb every hunk that overlaps or touches the group.
        loop {
            if i < ho.len() && ho[i].before.start <= hi {
                hi = hi.max(ho[i].before.end);
                i += 1;
            } else if j < ht.len() && ht[j].before.start <= hi {
                hi = hi.max(ht[j].before.end);
                j += 1;
            } else {
                break;
            }
        }
        for line in &b[pos as usize..lo as usize] {
            out.extend_from_slice(line);
        }
        let ours_hunks = &ho[gi..i];
        let theirs_hunks = &ht[gj..j];
        let side = |hs: &[Hunk], other: &[&[u8]]| -> Vec<u8> {
            let mut text = Vec::new();
            let mut p = lo;
            for h in hs {
                for line in &b[p as usize..h.before.start as usize] {
                    text.extend_from_slice(line);
                }
                for line in &other[h.after.start as usize..h.after.end as usize] {
                    text.extend_from_slice(line);
                }
                p = h.before.end;
            }
            for line in &b[p as usize..hi as usize] {
                text.extend_from_slice(line);
            }
            text
        };
        if theirs_hunks.is_empty() {
            out.extend_from_slice(&side(ours_hunks, &o));
        } else if ours_hunks.is_empty() {
            out.extend_from_slice(&side(theirs_hunks, &t));
        } else {
            let a = side(ours_hunks, &o);
            let c = side(theirs_hunks, &t);
            if a == c {
                out.extend_from_slice(&a);
            } else {
                conflicts += 1;
                if !out.is_empty() && !out.ends_with(b"\n") {
                    out.push(b'\n');
                }
                out.extend_from_slice(format!("<<<<<<< {ours_label}\n").as_bytes());
                out.extend_from_slice(&a);
                if !a.is_empty() && !a.ends_with(b"\n") {
                    out.push(b'\n');
                }
                out.extend_from_slice(b"=======\n");
                out.extend_from_slice(&c);
                if !c.is_empty() && !c.ends_with(b"\n") {
                    out.push(b'\n');
                }
                out.extend_from_slice(format!(">>>>>>> {theirs_label}\n").as_bytes());
            }
        }
        pos = hi;
    }
    for line in &b[pos as usize..] {
        out.extend_from_slice(line);
    }
    TextMerge { content: out, conflicts }
}

/// Merges JSON objects key by key. None when the files aren't all JSON
/// objects or both sides changed the same key differently.
pub fn json(base: Option<&[u8]>, ours: &[u8], theirs: &[u8]) -> Option<Vec<u8>> {
    let parse = |b: &[u8]| serde_json::from_slice::<Value>(b).ok();
    let o = parse(ours)?;
    let t = parse(theirs)?;
    let b = match base {
        Some(b) => parse(b)?,
        None => Value::Object(Default::default()),
    };
    if !(o.is_object() && t.is_object() && b.is_object()) {
        return None;
    }
    let merged = json_value(Some(&b), Some(&o), Some(&t))??;
    // Write it back in ours' style.
    let text = String::from_utf8_lossy(ours);
    let indent = text
        .lines()
        .skip(1)
        .find(|l| l.starts_with([' ', '\t']))
        .map(|l| l.chars().take_while(|c| *c == ' ' || *c == '\t').collect::<String>())
        .unwrap_or_else(|| "  ".into());
    let mut out = Vec::new();
    let formatter = serde_json::ser::PrettyFormatter::with_indent(indent.as_bytes());
    let mut ser = serde_json::Serializer::with_formatter(&mut out, formatter);
    serde::Serialize::serialize(&merged, &mut ser).ok()?;
    if ours.ends_with(b"\n") {
        out.push(b'\n');
    }
    Some(out)
}

/// Some(None): the key is gone. None: a conflict.
fn json_value(b: Option<&Value>, o: Option<&Value>, t: Option<&Value>) -> Option<Option<Value>> {
    if o == t {
        return Some(o.cloned());
    }
    if b == o {
        return Some(t.cloned());
    }
    if b == t {
        return Some(o.cloned());
    }
    let (Some(Value::Object(om)), Some(Value::Object(tm))) = (o, t) else { return None };
    let empty = serde_json::Map::new();
    let bm = match b {
        Some(Value::Object(m)) => m,
        None => &empty,
        _ => return None,
    };
    let mut out = serde_json::Map::new();
    for (k, ov) in om {
        if let Some(v) = json_value(bm.get(k), Some(ov), tm.get(k))? {
            out.insert(k.clone(), v);
        }
    }
    // Keys only theirs has, placed after the key they follow in theirs.
    let keys: Vec<&String> = tm.keys().collect();
    for (idx, k) in keys.iter().enumerate() {
        if om.contains_key(*k) || bm.contains_key(*k) {
            if !om.contains_key(*k) {
                // Deleted on ours: handled above only for keys ours has.
                if let Some(v) = json_value(bm.get(*k), None, tm.get(*k))? {
                    out.insert((*k).clone(), v);
                }
            }
            continue;
        }
        let value = tm[*k].clone();
        let after = keys[..idx].iter().rev().find(|p| out.contains_key(**p)).map(|p| (*p).clone());
        insert_after(&mut out, after.as_deref(), (*k).clone(), value);
    }
    Some(Some(Value::Object(out)))
}

fn insert_after(map: &mut serde_json::Map<String, Value>, after: Option<&str>, key: String, value: Value) {
    let Some(after) = after else {
        // Theirs put it first.
        let rest = std::mem::take(map);
        map.insert(key, value);
        map.extend(rest);
        return;
    };
    let rest: Vec<(String, Value)> = std::mem::take(map).into_iter().collect();
    for (k, v) in rest {
        let hit = k == after;
        map.insert(k, v);
        if hit {
            map.insert(key.clone(), value.clone());
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    /// Landing on the server: any conflict stops it.
    Land,
    /// Syncing a workspace: conflicts are written into files.
    Sync,
}

/// A conflict the merge found.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Conflict {
    pub path: String,
    /// text, binary, deleted, type.
    pub kind: String,
}

/// Something sync resolved on its own that the person should know about.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Note {
    pub path: String,
    pub message: String,
}

pub struct TreeMerge {
    pub tree: Id,
    /// Objects the merge created (merged blobs and trees).
    pub objects: MemSource,
    pub conflicts: Vec<Conflict>,
    pub notes: Vec<Note>,
}

impl TreeMerge {
    pub fn clean(&self) -> bool {
        self.conflicts.is_empty()
    }
}

/// Merges trees: `base` is where the change started, `ours` is main and
/// `theirs` is the change.
pub fn trees(src: &dyn Source, base: Option<&Id>, ours: &Id, theirs: &Id, mode: Mode, labels: (&str, &str)) -> Result<TreeMerge> {
    let mut m = Merger { src, mode, labels, objects: MemSource::default(), conflicts: Vec::new(), notes: Vec::new() };
    let empty = tree::empty_id();
    m.objects.put(tree::empty());
    let base = base.copied().unwrap_or(empty);
    let tree = m.merge_tree(&base, ours, theirs, "")?;
    Ok(TreeMerge { tree, objects: m.objects, conflicts: m.conflicts, notes: m.notes })
}

struct Merger<'a> {
    src: &'a dyn Source,
    mode: Mode,
    labels: (&'a str, &'a str),
    objects: MemSource,
    conflicts: Vec<Conflict>,
    notes: Vec<Note>,
}

impl Merger<'_> {
    fn read(&self, id: &Id) -> Result<Vec<TreeEntry>> {
        let src = Layered { first: &self.objects, second: self.src };
        tree::read(&src, id)
    }

    fn file(&self, id: &Id) -> Result<Object> {
        let src = Layered { first: &self.objects, second: self.src };
        src.must(id)
    }

    fn merge_tree(&mut self, base: &Id, ours: &Id, theirs: &Id, prefix: &str) -> Result<Id> {
        if ours == theirs || base == theirs {
            return Ok(*ours);
        }
        if base == ours {
            return Ok(*theirs);
        }
        let map = |es: Vec<TreeEntry>| es.into_iter().map(|e| (e.name.clone(), e)).collect::<BTreeMap<_, _>>();
        let b = map(self.read(base)?);
        let o = map(self.read(ours)?);
        let t = map(self.read(theirs)?);
        let mut names: Vec<&String> = b.keys().chain(o.keys()).chain(t.keys()).collect();
        names.sort();
        names.dedup();
        let mut out = Vec::new();
        for name in names {
            let path = if prefix.is_empty() { name.clone() } else { format!("{prefix}/{name}") };
            if let Some(e) = self.merge_entry(b.get(name), o.get(name), t.get(name), &path)? {
                out.push(TreeEntry { name: name.clone(), kind: e.0, id: e.1 });
            }
        }
        Ok(self.objects.put(Object::new(Kind::Tree, tree::encode(&out))))
    }

    fn merge_entry(&mut self, b: Option<&TreeEntry>, o: Option<&TreeEntry>, t: Option<&TreeEntry>, path: &str) -> Result<Option<(EntryKind, Id)>> {
        let key = |e: Option<&TreeEntry>| e.map(|e| (e.kind, e.id));
        let (kb, ko, kt) = (key(b), key(o), key(t));
        if ko == kt || kb == kt {
            return Ok(ko);
        }
        if kb == ko {
            return Ok(kt);
        }
        let is_tree = |e: Option<&TreeEntry>| e.is_some_and(|e| e.kind == EntryKind::Tree);
        let is_file = |e: Option<&TreeEntry>| e.is_some_and(|e| e.kind != EntryKind::Tree);
        match (o, t) {
            (Some(oe), Some(te)) if is_tree(o) && is_tree(t) => {
                let base = if is_tree(b) { b.unwrap().id } else { tree::empty_id() };
                Ok(Some((EntryKind::Tree, self.merge_tree(&base, &oe.id, &te.id, path)?)))
            }
            (Some(oe), Some(te)) if is_file(o) && is_file(t) => self.merge_file(b.filter(|e| e.kind != EntryKind::Tree), oe, te, path),
            (None, Some(te)) => {
                // Main deleted what the change edited.
                self.unmarkable(path, "deleted", Some(te), format!("main deleted {path}; kept your version"))
            }
            (Some(_), None) => {
                // The change deleted what main edited.
                self.unmarkable(path, "deleted", None, format!("you deleted {path}, which main also changed; `vcs restore {path} --from main` brings main's back"))
            }
            _ => {
                // A file on one side, a folder on the other.
                self.unmarkable(path, "type", t, format!("{path} is a file on one side and a folder on the other; kept yours"))
            }
        }
    }

    /// A conflict that can't be written with markers: landing stops; sync
    /// keeps the change's side (`keep`) and says so.
    fn unmarkable(&mut self, path: &str, kind: &'static str, keep: Option<&TreeEntry>, note: String) -> Result<Option<(EntryKind, Id)>> {
        self.conflicts.push(Conflict { path: path.to_string(), kind: kind.to_string() });
        if self.mode == Mode::Sync {
            self.conflicts.pop();
            self.notes.push(Note { path: path.to_string(), message: note });
        }
        Ok(keep.map(|e| (e.kind, e.id)))
    }

    fn merge_file(&mut self, b: Option<&TreeEntry>, o: &TreeEntry, t: &TreeEntry, path: &str) -> Result<Option<(EntryKind, Id)>> {
        // The kind (exec bit, link) merges on its own.
        let kind = if o.kind == t.kind {
            o.kind
        } else if b.is_some_and(|b| b.kind == o.kind) {
            t.kind
        } else if b.is_some_and(|b| b.kind == t.kind) {
            o.kind
        } else {
            t.kind
        };
        if o.id == t.id {
            return Ok(Some((kind, o.id)));
        }
        if b.is_some_and(|b| b.id == o.id) {
            return Ok(Some((kind, t.id)));
        }
        if b.is_some_and(|b| b.id == t.id) {
            return Ok(Some((kind, o.id)));
        }
        let oo = self.file(&o.id)?;
        let to = self.file(&t.id)?;
        let bo = match b {
            Some(b) => Some(self.file(&b.id)?),
            None => None,
        };
        let text = |obj: &Object| obj.kind == Kind::Blob && !is_binary(&obj.data);
        let all_text = text(&oo) && text(&to) && bo.as_ref().is_none_or(text) && o.kind != EntryKind::Link && t.kind != EntryKind::Link;
        if !all_text {
            return self.unmarkable(path, "binary", Some(t), format!("{path} changed on main and in your change and can't be merged line by line; kept yours, `vcs restore {path} --from main` takes main's"));
        }
        let base_bytes: &[u8] = bo.as_ref().map(|b| b.data.as_slice()).unwrap_or(b"");
        let merged = text_merge(base_bytes, &oo.data, &to.data, self.labels.0, self.labels.1);
        if merged.conflicts == 0 {
            let id = self.put_file(merged.content);
            return Ok(Some((kind, id)));
        }
        if path.ends_with(".json")
            && let Some(content) = json(bo.as_ref().map(|b| b.data.as_slice()), &oo.data, &to.data)
        {
            let id = self.put_file(content);
            return Ok(Some((kind, id)));
        }
        self.conflicts.push(Conflict { path: path.to_string(), kind: "text".into() });
        match self.mode {
            Mode::Land => Ok(Some((kind, t.id))),
            Mode::Sync => {
                let id = self.put_file(merged.content);
                Ok(Some((kind, id)))
            }
        }
    }

    fn put_file(&mut self, content: Vec<u8>) -> Id {
        let objects = chunk::file_objects(&content);
        let id = objects[0].id();
        for o in objects {
            self.objects.put(o);
        }
        id
    }
}

fn text_merge(base: &[u8], ours: &[u8], theirs: &[u8], ol: &str, tl: &str) -> TextMerge {
    text(base, ours, theirs, ol, tl)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree::{Entry, Files, build, flatten};

    #[test]
    fn clean_text_merge() {
        let base = b"a\nb\nc\nd\ne\nf\n";
        let ours = b"A\nb\nc\nd\ne\nf\n";
        let theirs = b"a\nb\nc\nd\ne\nF\n";
        let m = text(base, ours, theirs, "main", "change");
        assert_eq!(m.conflicts, 0);
        assert_eq!(m.content, b"A\nb\nc\nd\ne\nF\n");
    }

    #[test]
    fn conflicting_text_merge() {
        let base = b"a\nb\nc\n";
        let ours = b"a\nB1\nc\n";
        let theirs = b"a\nB2\nc\n";
        let m = text(base, ours, theirs, "main 4", "your change bcdf");
        assert_eq!(m.conflicts, 1);
        assert_eq!(String::from_utf8(m.content).unwrap(), "a\n<<<<<<< main 4\nB1\n=======\nB2\n>>>>>>> your change bcdf\nc\n");
    }

    #[test]
    fn same_change_both_sides() {
        let m = text(b"a\nb\n", b"a\nX\n", b"a\nX\n", "o", "t");
        assert_eq!(m.conflicts, 0);
        assert_eq!(m.content, b"a\nX\n");
    }

    #[test]
    fn json_keys_merge() {
        let base = b"{\n  \"a\": \"1\"\n}\n";
        let ours = b"{\n  \"a\": \"1\",\n  \"b\": \"2\"\n}\n";
        let theirs = b"{\n  \"a\": \"1\",\n  \"c\": \"3\"\n}\n";
        // Line by line this is a conflict (both added after "a").
        assert_eq!(text(base, ours, theirs, "o", "t").conflicts, 1);
        let merged = json(Some(base), ours, theirs).unwrap();
        let v: Value = serde_json::from_slice(&merged).unwrap();
        assert_eq!(v["b"], "2");
        assert_eq!(v["c"], "3");
        assert!(json(Some(base), b"{\"a\": \"x\"}", b"{\"a\": \"y\"}").is_none());
    }

    fn put_files(src: &mut MemSource, list: &[(&str, &str)]) -> Id {
        let files: Files = list
            .iter()
            .map(|(p, c)| (p.to_string(), Entry { kind: EntryKind::File, id: src.put(Object::blob(c.as_bytes().to_vec())) }))
            .collect();
        let (root, trees) = build(&files);
        for t in trees {
            src.put(t);
        }
        root
    }

    #[test]
    fn tree_merges() {
        let mut src = MemSource::default();
        let base = put_files(&mut src, &[("a.txt", "a\nb\nc\n"), ("src/x.rs", "x\n"), ("keys.json", "{\n  \"a\": 1\n}\n")]);
        let ours = put_files(&mut src, &[("a.txt", "A\nb\nc\n"), ("src/x.rs", "x\n"), ("src/new.rs", "n\n"), ("keys.json", "{\n  \"a\": 1,\n  \"b\": 2\n}\n")]);
        let theirs = put_files(&mut src, &[("a.txt", "a\nb\nC\n"), ("src/x.rs", "x2\n"), ("keys.json", "{\n  \"a\": 1,\n  \"c\": 3\n}\n")]);
        let m = trees(&src, Some(&base), &ours, &theirs, Mode::Land, ("main", "change")).unwrap();
        assert!(m.clean(), "{:?}", m.conflicts);
        let layered = Layered { first: &m.objects, second: &src };
        let files = flatten(&layered, &m.tree).unwrap();
        let read = |p: &str| String::from_utf8(layered.must(&files[p].id).unwrap().data).unwrap();
        assert_eq!(read("a.txt"), "A\nb\nC\n");
        assert_eq!(read("src/x.rs"), "x2\n");
        assert_eq!(read("src/new.rs"), "n\n");
        let keys: Value = serde_json::from_str(&read("keys.json")).unwrap();
        assert_eq!((keys["b"].as_i64(), keys["c"].as_i64()), (Some(2), Some(3)));
    }

    #[test]
    fn tree_conflicts() {
        let mut src = MemSource::default();
        let base = put_files(&mut src, &[("a.txt", "a\n"), ("gone.txt", "g\n")]);
        let ours = put_files(&mut src, &[("a.txt", "main\n")]);
        let theirs = put_files(&mut src, &[("a.txt", "change\n"), ("gone.txt", "edited\n")]);
        let land = trees(&src, Some(&base), &ours, &theirs, Mode::Land, ("main", "change")).unwrap();
        let mut paths: Vec<_> = land.conflicts.iter().map(|c| (c.path.as_str(), c.kind.as_str())).collect();
        paths.sort();
        assert_eq!(paths, vec![("a.txt", "text"), ("gone.txt", "deleted")]);
        let sync = trees(&src, Some(&base), &ours, &theirs, Mode::Sync, ("main 2", "your change x")).unwrap();
        assert_eq!(sync.conflicts.len(), 1);
        assert_eq!(sync.notes.len(), 1);
        let layered = Layered { first: &sync.objects, second: &src };
        let files = flatten(&layered, &sync.tree).unwrap();
        let a = String::from_utf8(layered.must(&files["a.txt"].id).unwrap().data).unwrap();
        assert!(a.contains("<<<<<<< main 2"));
        assert!(files.contains_key("gone.txt"));
    }

    #[test]
    fn fast_paths() {
        let mut src = MemSource::default();
        let base = put_files(&mut src, &[("a.txt", "a\n")]);
        let theirs = put_files(&mut src, &[("a.txt", "b\n")]);
        let m = trees(&src, Some(&base), &base, &theirs, Mode::Land, ("main", "change")).unwrap();
        assert_eq!(m.tree, theirs);
    }
}