//! Smart merges for code: spots both sides changed where a line-by-line
//! merge sees a conflict but there's one clear answer.
//!
//! - **Both sides added lines at the same place** (two new functions, tests
//!   or changelog entries) and none of them in common: both go in, main's
//!   first.
//! - **Import blocks** (JS/TS, Python, Rust, Go, Java/Kotlin, C/C++, Swift,
//!   C#): each side's added imports go in and removed ones go, kept sorted
//!   when they were.
//! - **TOML** merges by table and key, as JSON does, keeping the file's
//!   layout and comments.

use std::collections::{BTreeMap, HashSet};

use crate::diff::lines;

#[derive(Clone, Copy, PartialEq, Eq)]
enum Lang {
    Js,
    Py,
    Rust,
    Go,
    Java,
    C,
    Swift,
    CSharp,
    Other,
}

fn lang(path: &str) -> Lang {
    let ext = path.rsplit_once('.').map(|(_, e)| e.to_ascii_lowercase()).unwrap_or_default();
    match ext.as_str() {
        "js" | "jsx" | "ts" | "tsx" | "mjs" | "cjs" | "mts" | "cts" => Lang::Js,
        "py" | "pyi" => Lang::Py,
        "rs" => Lang::Rust,
        "go" => Lang::Go,
        "java" | "kt" | "kts" | "scala" => Lang::Java,
        "c" | "h" | "cc" | "cpp" | "cxx" | "hpp" | "hh" | "m" | "mm" => Lang::C,
        "swift" => Lang::Swift,
        "cs" => Lang::CSharp,
        _ => Lang::Other,
    }
}

fn blank(line: &[u8]) -> bool {
    line.iter().all(|b| b.is_ascii_whitespace())
}

/// A whole import on one line, for the file's language.
fn is_import(lang: Lang, line: &[u8]) -> bool {
    let Ok(s) = std::str::from_utf8(line) else { return false };
    let s = s.trim();
    match lang {
        Lang::Js => {
            (s.starts_with("import ") && (s.contains(" from ") || s.starts_with("import \"") || s.starts_with("import '")))
                || (s.starts_with("export ") && s.contains(" from "))
        }
        Lang::Py => (s.starts_with("import ") || (s.starts_with("from ") && s.contains(" import "))) && !s.ends_with('('),
        Lang::Rust => (s.starts_with("use ") || s.starts_with("pub use ") || s.starts_with("pub(crate) use ")) && s.ends_with(';'),
        Lang::Go => (s.starts_with('"') && s.ends_with('"')) || (s.ends_with('"') && s.split_whitespace().count() == 2) || s.starts_with("import \""),
        Lang::Java => s.starts_with("import ") && (s.ends_with(';') || !s.contains(' ') || s.split_whitespace().count() == 2),
        Lang::C => s.starts_with("#include ") || s.starts_with("#import "),
        Lang::Swift => s.starts_with("import "),
        Lang::CSharp => s.starts_with("using ") && s.ends_with(';'),
        Lang::Other => false,
    }
}

/// Lines with a line break each (the file's last line may lack one).
fn ended(line: &[u8]) -> Vec<u8> {
    let mut l = line.to_vec();
    if !l.ends_with(b"\n") {
        l.push(b'\n');
    }
    l
}

/// The resolver for a file at `path` (see `merge::text_with`).
pub fn resolver(path: &str) -> impl Fn(&[&[u8]], &[u8], &[u8], (Option<&[u8]>, Option<&[u8]>)) -> Option<Vec<u8>> {
    let lang = lang(path);
    let path = path.to_string();
    move |base: &[&[u8]], ours: &[u8], theirs: &[u8], (before, after): (Option<&[u8]>, Option<&[u8]>)| {
        let o = lines(ours);
        let t = lines(theirs);
        if lang != Lang::Other {
            let all: Vec<&[u8]> = base.iter().copied().chain(o.iter().copied()).chain(t.iter().copied()).collect();
            if all.iter().any(|l| !blank(l)) && all.iter().all(|l| blank(l) || is_import(lang, l)) {
                // Two imports of one name (both sides changed `import { a }` from one module) don't merge.
                let added = |side: &[&[u8]]| -> HashSet<String> {
                    let kept: HashSet<&[u8]> = base.iter().copied().collect();
                    side.iter().filter(|l| !kept.contains(*l)).flat_map(|l| bound_names(lang, l)).collect()
                };
                if !added(&o).is_disjoint(&added(&t)) {
                    return None;
                }
                // Kept in order where the block around them is.
                let around = before.zip(after).is_some_and(|(b, a)| !blank(b) && !blank(a) && b.trim_ascii() <= a.trim_ascii());
                return Some(imports(base, &o, &t, around));
            }
        }
        if base.is_empty() && both_go_in(&path, lang, ours, theirs, &o, &t) {
            // Both added here: both go in, main's first, or in order where
            // they're entries of a sorted list (a Makefile's, #includes).
            let mut lines: Vec<&[u8]> = o.iter().chain(t.iter()).copied().collect();
            if sorted_list(&o, &t, before, after) {
                lines.sort_by_key(|l| l.trim_ascii());
            }
            let mut out = Vec::with_capacity(ours.len() + theirs.len() + 1);
            for l in lines {
                out.extend_from_slice(&ended(l));
            }
            return Some(out);
        }
        None
    }
}

/// Lines each side added one per entry, in order, between neighbours in
/// order too: a sorted list both sides added to.
fn sorted_list(o: &[&[u8]], t: &[&[u8]], before: Option<&[u8]>, after: Option<&[u8]>) -> bool {
    let key = |l: &[u8]| l.trim_ascii().to_vec();
    let added: Vec<&[u8]> = o.iter().chain(t.iter()).copied().collect();
    if added.iter().any(|l| blank(l)) || added.len() < 2 {
        return false;
    }
    let ordered = |side: &[&[u8]]| side.windows(2).all(|w| key(w[0]) <= key(w[1]));
    let (Some(b), Some(a)) = (before, after) else { return false };
    if blank(b) || blank(a) || key(b) > key(a) {
        return false;
    }
    ordered(o) && ordered(t) && added.iter().all(|l| key(b) <= key(l) && key(l) <= key(a))
}

/// Whether two blocks added at one spot can both go in: they share no lines
/// (closing braces aside), declare nothing in common, and in code each is
/// whole (brackets closed) and made of declarations or test cases, not
/// statements whose order matters.
fn both_go_in(path: &str, lang: Lang, ours: &[u8], theirs: &[u8], o: &[&[u8]], t: &[&[u8]]) -> bool {
    let closing = |l: &[u8]| l.iter().all(|c| c.is_ascii_whitespace() || b"}])>;,".contains(c));
    let mine: HashSet<&[u8]> = o.iter().copied().filter(|l| !blank(l) && !closing(l)).collect();
    if t.iter().any(|l| !blank(l) && !closing(l) && mine.contains(l)) {
        return false;
    }
    let names = |b: &[u8]| crate::decls::declared(path, b).into_iter().map(|(k, _)| k).collect::<HashSet<_>>();
    if !names(ours).is_disjoint(&names(theirs)) {
        return false;
    }
    if lang == Lang::Other {
        return true;
    }
    [ours, theirs].iter().all(|block| balanced(block) && whole_items(path, block))
}

/// Brackets open and close within the block (strings and comments aside, roughly).
fn balanced(block: &[u8]) -> bool {
    let mut depth = [0i64; 3];
    let (mut quote, mut escaped) = (None::<u8>, false);
    for line in lines(block) {
        let line = match line.windows(2).position(|w| w == b"//") {
            Some(p) if quote.is_none() => &line[..p],
            _ => line,
        };
        for &c in line {
            if let Some(q) = quote {
                if escaped {
                    escaped = false;
                } else if c == b'\\' {
                    escaped = true;
                } else if c == q {
                    quote = None;
                }
                continue;
            }
            match c {
                b'"' | b'`' => quote = Some(c),
                b'(' => depth[0] += 1,
                b')' => depth[0] -= 1,
                b'[' => depth[1] += 1,
                b']' => depth[1] -= 1,
                b'{' => depth[2] += 1,
                b'}' => depth[2] -= 1,
                _ => {}
            }
            if depth.iter().any(|d| *d < 0) {
                return false;
            }
        }
        if quote == Some(b'"') {
            quote = None;
        }
    }
    depth == [0, 0, 0] && quote.is_none()
}

/// A block made of whole declarations (functions, classes, types, constants)
/// or test cases (`it("…", …)`, `test(…)`), with comments, decorators and
/// blank lines between: things whose order next to each other doesn't matter.
fn whole_items(path: &str, block: &[u8]) -> bool {
    let decl_lines: HashSet<usize> = crate::decls::declared(path, block).into_iter().map(|(_, n)| n).collect();
    let mut seen_item = false;
    let mut depth = 0i64;
    for (n, line) in String::from_utf8_lossy(block).lines().enumerate() {
        let s = line.trim();
        let top = depth == 0;
        depth += s.matches(['{', '(', '[']).count() as i64 - s.matches(['}', ')', ']']).count() as i64;
        if !top || s.is_empty() || s.starts_with("//") || s.starts_with('#') && !s.starts_with("#[") || s.starts_with("/*") || s.starts_with('*') || s.starts_with('@') || s.starts_with("#[") {
            continue;
        }
        let test_case = ["it(", "test(", "describe(", "bench(", "it.each", "test.each", "#[test]"].iter().any(|p| s.starts_with(p));
        if decl_lines.contains(&n) || test_case {
            seen_item = true;
            continue;
        }
        // Python bodies are indented; a top-level line that isn't an item is a statement.
        if line.starts_with(char::is_whitespace) && seen_item {
            continue;
        }
        return false;
    }
    seen_item
}

/// The names an import line binds (`import { a, b as c } from "x"` → a, c).
fn bound_names(lang: Lang, line: &[u8]) -> Vec<String> {
    let Ok(s) = std::str::from_utf8(line) else { return Vec::new() };
    let s = s.trim().trim_end_matches(';');
    let split = |list: &str| -> Vec<String> {
        list.split(',')
            .map(|p| p.trim())
            .filter(|p| !p.is_empty() && *p != "*" && *p != "self")
            .map(|p| p.rsplit(" as ").next().unwrap_or(p).trim().trim_start_matches("type ").to_string())
            .collect()
    };
    match lang {
        Lang::Js => {
            let Some(clause) = s.strip_prefix("import ").or_else(|| s.strip_prefix("export ")) else { return Vec::new() };
            let Some((clause, _)) = clause.split_once(" from ") else { return Vec::new() };
            let clause = clause.trim().trim_start_matches("type ");
            let mut names = Vec::new();
            let (head, braces) = match clause.find('{') {
                Some(i) => (&clause[..i], Some(clause[i + 1..].trim_end_matches('}'))),
                None => (clause, None),
            };
            for part in head.split(',').map(str::trim).filter(|p| !p.is_empty()) {
                names.push(part.rsplit(" as ").next().unwrap_or(part).trim().to_string());
            }
            if let Some(b) = braces {
                names.extend(split(b));
            }
            names
        }
        Lang::Py => match s.split_once(" import ") {
            Some((_, list)) if s.starts_with("from ") => split(list.trim_matches(['(', ')'])),
            _ => s.strip_prefix("import ").map(split).unwrap_or_default().into_iter().map(|n| n.split('.').next().unwrap_or(&n).to_string()).collect(),
        },
        Lang::Rust => {
            let path = s.trim_start_matches("pub(crate) ").trim_start_matches("pub ").trim_start_matches("use ").trim();
            match path.find('{') {
                Some(i) => split(path[i + 1..].trim_end_matches('}')),
                None => vec![path.rsplit(" as ").next().unwrap_or(path).rsplit("::").next().unwrap_or(path).to_string()],
            }
        }
        Lang::Go => {
            let parts: Vec<&str> = s.trim_start_matches("import ").split_whitespace().collect();
            match parts.as_slice() {
                [alias, _] => vec![alias.to_string()],
                [path] => vec![path.trim_matches('"').rsplit('/').next().unwrap_or("").to_string()],
                _ => Vec::new(),
            }
        }
        // Repeated includes and usings are harmless there.
        Lang::Java | Lang::C | Lang::Swift | Lang::CSharp | Lang::Other => Vec::new(),
    }
}

/// An import block: what either side added, less what either side removed.
/// `around`: the lines just around the spot are in order.
fn imports(base: &[&[u8]], o: &[&[u8]], t: &[&[u8]], around: bool) -> Vec<u8> {
    let key = |l: &[u8]| String::from_utf8_lossy(l).trim().to_string();
    let in_o: HashSet<String> = o.iter().map(|l| key(l)).collect();
    let in_t: HashSet<String> = t.iter().map(|l| key(l)).collect();
    let removed: HashSet<String> = base.iter().map(|l| key(l)).filter(|k| !k.is_empty() && (!in_o.contains(k) || !in_t.contains(k))).collect();
    let mut seen: HashSet<String> = HashSet::new();
    let mut merged: Vec<Vec<u8>> = Vec::new();
    for l in o.iter().chain(t.iter()) {
        let k = key(l);
        if removed.contains(&k) || (!k.is_empty() && !seen.insert(k)) {
            continue;
        }
        merged.push(ended(l));
    }
    let sorted = |ls: &[&[u8]]| ls.len() > 1 && ls.windows(2).all(|w| key(w[0]) <= key(w[1]));
    if !merged.iter().any(|l| blank(l)) && (sorted(base) || sorted(o) || (around && (sorted(t) || t.len() < 2) && (sorted(o) || o.len() < 2))) {
        merged.sort_by_key(|l| key(l));
    }
    merged.concat()
}

/// A TOML file as tables and keys: each key's lines (a value can run over
/// several), and which table each line is in.
struct Toml {
    lines: Vec<String>,
    /// (table, key) → the line range of the key and its value.
    keys: BTreeMap<(String, String), (usize, usize)>,
    tables: Vec<String>,
}

fn parse_toml(text: &str) -> Option<Toml> {
    let lines: Vec<String> = text.split_inclusive('\n').map(String::from).collect();
    let mut keys = BTreeMap::new();
    let mut tables = vec![String::new()];
    let mut table = String::new();
    let mut i = 0;
    while i < lines.len() {
        let s = lines[i].trim();
        if s.starts_with("[[") {
            return None;
        }
        if s.starts_with('[') {
            table = s.trim_start_matches('[').split(']').next()?.trim().to_string();
            tables.push(table.clone());
            i += 1;
            continue;
        }
        if s.is_empty() || s.starts_with('#') {
            i += 1;
            continue;
        }
        let (key, value) = s.split_once('=')?;
        let key = key.trim().to_string();
        // A value runs on while its brackets are open or a triple-quoted string is.
        let mut depth: i64 = 0;
        let mut quoted = false;
        let mut end = i;
        let mut text = value.to_string();
        loop {
            for q in ["\"\"\"", "'''"] {
                if text.matches(q).count() % 2 == 1 {
                    quoted = !quoted;
                }
            }
            if !quoted {
                depth += text.matches(['[', '{']).count() as i64 - text.matches([']', '}']).count() as i64;
            }
            if (depth <= 0 && !quoted) || end + 1 >= lines.len() {
                break;
            }
            end += 1;
            text = lines[end].clone();
        }
        if keys.insert((table.clone(), key), (i, end + 1)).is_some() {
            return None;
        }
        i = end + 1;
    }
    Some(Toml { lines, keys, tables })
}

/// Merges TOML files by table and key, from main's file: None when both
/// sides changed the same key differently (or it isn't TOML this reads).
pub fn toml(base: Option<&[u8]>, ours: &[u8], theirs: &[u8]) -> Option<Vec<u8>> {
    let read = |b: &[u8]| std::str::from_utf8(b).ok().and_then(parse_toml);
    let b = match base {
        Some(b) => read(b)?,
        None => parse_toml("")?,
    };
    let o = read(ours)?;
    let t = read(theirs)?;
    let text = |doc: &Toml, k: &(String, String)| doc.keys.get(k).map(|(s, e)| doc.lines[*s..*e].concat());
    let all: Vec<(String, String)> = b.keys.keys().chain(o.keys.keys()).chain(t.keys.keys()).cloned().collect::<std::collections::BTreeSet<_>>().into_iter().collect();
    // What to do to main's file: replace a key's lines, drop them, or add a key to a table.
    let mut replace: BTreeMap<usize, (usize, Option<String>)> = BTreeMap::new();
    let mut added: Vec<((String, String), String)> = Vec::new();
    for k in &all {
        let (vb, vo, vt) = (text(&b, k), text(&o, k), text(&t, k));
        if vo == vt || vb == vt {
            continue;
        }
        if vb != vo {
            return None;
        }
        match (o.keys.get(k), vt) {
            (Some((s, e)), next) => {
                replace.insert(*s, (*e, next));
            }
            (None, Some(next)) => added.push((k.clone(), next)),
            (None, None) => {}
        }
    }
    let mut out: Vec<String> = Vec::new();
    let mut table = String::new();
    // After a table's last key in main's file: the keys the change added to it.
    let last_line: BTreeMap<String, usize> = o.keys.iter().fold(BTreeMap::new(), |mut m, ((tb, _), (_, e))| {
        let v = m.entry(tb.clone()).or_insert(0);
        *v = (*v).max(*e);
        m
    });
    let mut placed: HashSet<usize> = HashSet::new();
    let mut i = 0;
    while i < o.lines.len() {
        let s = o.lines[i].trim();
        if s.starts_with('[') {
            table = s.trim_start_matches('[').split(']').next().unwrap_or("").trim().to_string();
        }
        if let Some((end, next)) = replace.get(&i) {
            if let Some(n) = next {
                out.push(n.clone());
            }
            i = *end;
        } else {
            out.push(o.lines[i].clone());
            i += 1;
        }
        let at = i;
        if last_line.get(&table) == Some(&at) {
            for (n, ((tb, _), v)) in added.iter().enumerate() {
                if *tb == table && placed.insert(n) {
                    if out.last().is_some_and(|l| !l.ends_with('\n')) {
                        out.push("\n".into());
                    }
                    out.push(v.clone());
                }
            }
        }
    }
    // Keys for tables main's file has no keys in (or doesn't have): at the end, under their table.
    for (n, ((tb, _), v)) in added.iter().enumerate() {
        if placed.contains(&n) {
            continue;
        }
        if out.last().is_some_and(|l| !l.ends_with('\n')) {
            out.push("\n".into());
        }
        if !tb.is_empty() && !o.tables.contains(tb) && !out.iter().any(|l| l.trim() == format!("[{tb}]")) {
            out.push(format!("\n[{tb}]\n"));
        }
        out.push(v.clone());
    }
    let _ = &t.tables;
    Some(out.concat().into_bytes())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::merge;

    fn smart(path: &str, base: &str, ours: &str, theirs: &str) -> (String, usize) {
        let m = merge::text_for(path, base.as_bytes(), ours.as_bytes(), theirs.as_bytes(), "main", "change");
        (String::from_utf8(m.content).unwrap(), m.conflicts)
    }

    #[test]
    fn both_added_at_the_same_place() {
        let base = "fn a() {}\n";
        let ours = "fn a() {}\nfn b() {}\n";
        let theirs = "fn a() {}\nfn c() {}\n";
        assert_eq!(smart("lib.rs", base, ours, theirs), ("fn a() {}\nfn b() {}\nfn c() {}\n".into(), 0));
        // Some of the same lines: still a conflict.
        assert_eq!(smart("lib.rs", base, "fn a() {}\nfn b() {}\n}\n", "fn a() {}\nfn c() {}\n}\n").1, 1);
    }

    #[test]
    fn import_blocks() {
        let base = "import { a } from \"a\";\nimport { c } from \"c\";\n\nrun();\n";
        let ours = "import { a } from \"a\";\nimport { b } from \"b\";\nimport { c } from \"c\";\n\nrun();\n";
        let theirs = "import { a } from \"a\";\nimport { bb } from \"bb\";\nimport { c } from \"c\";\n\nrun();\n";
        let (text, conflicts) = smart("app.ts", base, ours, theirs);
        assert_eq!(conflicts, 0);
        assert_eq!(text, "import { a } from \"a\";\nimport { b } from \"b\";\nimport { bb } from \"bb\";\nimport { c } from \"c\";\n\nrun();\n");
        // One side removes an import the other kept.
        let (text, conflicts) = smart("m.py", "import os\nimport re\n", "import os\nimport sys\n", "import json\nimport os\nimport re\n");
        assert_eq!((text.as_str(), conflicts), ("import json\nimport os\nimport sys\n", 0));
    }

    #[test]
    fn toml_by_key() {
        let base = "[package]\nname = \"app\"\nversion = \"0.1.0\"\n\n# deps\n[dependencies]\nserde = \"1\"\n";
        let ours = "[package]\nname = \"app\"\nversion = \"0.2.0\"\n\n# deps\n[dependencies]\nserde = \"1\"\ntokio = \"1\"\n";
        let theirs = "[package]\nname = \"app\"\nversion = \"0.1.0\"\nedition = \"2024\"\n\n# deps\n[dependencies]\nserde = \"1\"\nrand = \"0.8\"\n";
        let merged = String::from_utf8(toml(Some(base.as_bytes()), ours.as_bytes(), theirs.as_bytes()).unwrap()).unwrap();
        assert_eq!(merged, "[package]\nname = \"app\"\nversion = \"0.2.0\"\nedition = \"2024\"\n\n# deps\n[dependencies]\nserde = \"1\"\ntokio = \"1\"\nrand = \"0.8\"\n");
        // The same key changed differently: no merge.
        assert!(toml(Some(base.as_bytes()), base.replace("0.1.0", "0.3.0").as_bytes(), base.replace("0.1.0", "0.4.0").as_bytes()).is_none());
    }
}