//! 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<Vec<u8>> {
    let lang = lang(path);
    move |base: &[&[u8]], ours: &[u8], theirs: &[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)) {
                return Some(imports(base, &o, &t));
            }
        }
        if base.is_empty() {
            // Both added here: both go in, unless they added some of the same lines.
            let mine: HashSet<&[u8]> = o.iter().copied().filter(|l| !blank(l)).collect();
            if t.iter().any(|l| !blank(l) && mine.contains(l)) {
                return None;
            }
            let mut out = Vec::with_capacity(ours.len() + theirs.len() + 1);
            for l in o.iter().chain(t.iter()) {
                out.extend_from_slice(&ended(l));
            }
            return Some(out);
        }
        None
    }
}

/// An import block: what either side added, less what either side removed.
fn imports(base: &[&[u8]], o: &[&[u8]], t: &[&[u8]]) -> 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)) {
        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());
    }
}