//! Syntax highlighting as ranges, not markup: the page carries the text
//! once, and a few lines of script hand these ranges to the CSS Custom
//! Highlight API. Offsets are in UTF-16 code units, as the DOM counts.
//!
//! Code is parsed with tree-sitter, whose grammars take tens of megabytes a
//! second; a language without one here is shown plain.

use crate::diff::{Kind, Row};
use std::cell::RefCell;
use std::fmt::Write;
use std::sync::LazyLock;
use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};

/// Larger files are shown plain: highlighting them costs more than it gives.
pub const LIMIT: usize = 1024 * 1024;

/// The highlight names for code, each styled by `::highlight(name)`.
const CODE: &[&str] = &[
    "comment", "string", "keyword", "number", "function", "type", "tag", "attr",
];

/// The captures the grammars' queries use, and the highlight each gives.
/// A capture takes its longest listed prefix: `function.method` is a
/// `function`.
const CAPTURES: [(&str, usize); 23] = [
    ("comment", 0),
    ("string", 1),
    ("string.special.key", 7),
    ("escape", 1),
    ("character", 1),
    ("keyword", 2),
    ("variable.builtin", 2),
    ("number", 3),
    ("float", 3),
    ("boolean", 3),
    ("constant", 3),
    ("constant.builtin", 3),
    ("function", 4),
    ("function.builtin", 4),
    ("function.method", 4),
    ("function.macro", 4),
    ("type", 5),
    ("type.builtin", 5),
    ("constructor", 5),
    ("tag", 6),
    ("attribute", 7),
    ("property.definition", 7),
    ("label", 2),
];

/// The names for a diff's rows: the code names, then the changed part of
/// an added line, and of a removed one.
const ROWS: &[&str] = &[
    "comment", "string", "keyword", "number", "function", "type", "tag", "attr", "xins", "xdel",
];

struct Language {
    name: &'static str,
    extensions: &'static [&'static str],
    config: HighlightConfiguration,
}

fn language(
    name: &'static str,
    extensions: &'static [&'static str],
    grammar: tree_sitter_language::LanguageFn,
    highlights: &[&str],
    injections: &str,
) -> Language {
    let mut config =
        HighlightConfiguration::new(grammar.into(), name, &highlights.join("\n"), injections, "")
            .expect("a bundled highlight query compiles");
    let names: Vec<&str> = CAPTURES.iter().map(|(n, _)| *n).collect();
    config.configure(&names);
    Language {
        name,
        extensions,
        config,
    }
}

static LANGUAGES: LazyLock<Vec<Language>> = LazyLock::new(|| {
    use tree_sitter_javascript as js;
    use tree_sitter_typescript as ts;
    vec![
        language(
            "rust",
            &["rs"],
            tree_sitter_rust::LANGUAGE,
            &[tree_sitter_rust::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "javascript",
            &["js", "mjs", "cjs", "jsx"],
            js::LANGUAGE,
            &[js::HIGHLIGHT_QUERY, js::JSX_HIGHLIGHT_QUERY],
            "",
        ),
        language(
            "typescript",
            &["ts", "mts", "cts"],
            ts::LANGUAGE_TYPESCRIPT,
            &[js::HIGHLIGHT_QUERY, ts::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "tsx",
            &["tsx"],
            ts::LANGUAGE_TSX,
            &[
                js::HIGHLIGHT_QUERY,
                js::JSX_HIGHLIGHT_QUERY,
                ts::HIGHLIGHTS_QUERY,
            ],
            "",
        ),
        language(
            "python",
            &["py", "pyi"],
            tree_sitter_python::LANGUAGE,
            &[tree_sitter_python::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "go",
            &["go"],
            tree_sitter_go::LANGUAGE,
            &[tree_sitter_go::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "c",
            &["c", "h"],
            tree_sitter_c::LANGUAGE,
            &[tree_sitter_c::HIGHLIGHT_QUERY],
            "",
        ),
        language(
            "cpp",
            &["cc", "cpp", "cxx", "hpp", "hh", "hxx"],
            tree_sitter_cpp::LANGUAGE,
            &[
                tree_sitter_c::HIGHLIGHT_QUERY,
                tree_sitter_cpp::HIGHLIGHT_QUERY,
            ],
            "",
        ),
        language(
            "java",
            &["java"],
            tree_sitter_java::LANGUAGE,
            &[tree_sitter_java::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "ruby",
            &["rb", "rake", "gemspec", "Gemfile", "Rakefile"],
            tree_sitter_ruby::LANGUAGE,
            &[tree_sitter_ruby::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "bash",
            &["sh", "bash", "zsh"],
            tree_sitter_bash::LANGUAGE,
            &[tree_sitter_bash::HIGHLIGHT_QUERY],
            "",
        ),
        language(
            "json",
            &["json"],
            tree_sitter_json::LANGUAGE,
            // Where patterns overlap the last wins: keys over strings.
            &[
                tree_sitter_json::HIGHLIGHTS_QUERY,
                "(pair key: (_) @string.special.key)",
            ],
            "",
        ),
        language(
            "toml",
            &["toml", "lock"],
            tree_sitter_toml_ng::LANGUAGE,
            &[tree_sitter_toml_ng::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "yaml",
            &["yml", "yaml"],
            tree_sitter_yaml::LANGUAGE,
            &[tree_sitter_yaml::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "html",
            &["html", "htm"],
            tree_sitter_html::LANGUAGE,
            &[tree_sitter_html::HIGHLIGHTS_QUERY],
            tree_sitter_html::INJECTIONS_QUERY,
        ),
        language(
            "css",
            &["css"],
            tree_sitter_css::LANGUAGE,
            &[tree_sitter_css::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "lua",
            &["lua"],
            tree_sitter_lua::LANGUAGE,
            &[tree_sitter_lua::HIGHLIGHTS_QUERY],
            "",
        ),
        language(
            "zig",
            &["zig"],
            tree_sitter_zig::LANGUAGE,
            &[tree_sitter_zig::HIGHLIGHTS_QUERY],
            "",
        ),
    ]
});

thread_local! {
    static HIGHLIGHTER: RefCell<Highlighter> = RefCell::new(Highlighter::new());
}

fn by_name(name: &str) -> Option<&'static Language> {
    LANGUAGES.iter().find(|l| l.name == name)
}

/// The language of the file at `path`: by its extension, or its name, or
/// the interpreter its first line names.
fn language_for(path: &str, text: &str) -> Option<&'static Language> {
    let name = path.rsplit('/').next().unwrap_or(path);
    let ext = name.rsplit_once('.').map_or(name, |(_, ext)| ext);
    let by_ext = LANGUAGES.iter().find(|l| l.extensions.contains(&ext));
    let shebang = text
        .lines()
        .next()
        .and_then(|l| l.strip_prefix("#!"))
        .unwrap_or_default();
    let interpreter = shebang.rsplit(['/', ' ']).next().unwrap_or_default();
    let by_shebang = match interpreter.trim_end_matches(|c: char| c.is_ascii_digit() || c == '.') {
        "sh" | "bash" | "zsh" => by_name("bash"),
        "python" => by_name("python"),
        "node" => by_name("javascript"),
        "ruby" => by_name("ruby"),
        _ => None,
    };
    by_ext.or(by_shebang)
}

/// Runs of each named highlight, as `[start, end)` in UTF-16 units.
#[derive(Debug, PartialEq)]
pub struct Ranges {
    names: &'static [&'static str],
    runs: Vec<Vec<(u32, u32)>>,
}

impl Ranges {
    fn new(names: &'static [&'static str]) -> Ranges {
        Ranges {
            names,
            runs: vec![Vec::new(); names.len()],
        }
    }

    /// Adds a run, joining it to the one before when they touch.
    fn push(&mut self, class: usize, start: u32, end: u32) {
        let runs = &mut self.runs[class];
        match runs.last_mut() {
            Some(last) if last.1 == start => last.1 = end,
            _ if start < end => runs.push((start, end)),
            _ => {}
        }
    }

    /// `name:gap,len,gap,len;…`, each gap from the end of the run before:
    /// short, and safe in an attribute as it is.
    pub fn encode(&self) -> String {
        let mut out = String::new();
        let named = self
            .names
            .iter()
            .zip(&self.runs)
            .filter(|(_, r)| !r.is_empty());
        for (name, runs) in named {
            let sep = if out.is_empty() { "" } else { ";" };
            write!(out, "{sep}{name}:").unwrap();
            let mut at = 0;
            for (i, (start, end)) in runs.iter().enumerate() {
                let sep = if i == 0 { "" } else { "," };
                write!(out, "{sep}{},{}", start - at, end - start).unwrap();
                at = *end;
            }
        }
        out
    }
}

fn utf16_len(s: &str) -> u32 {
    s.chars().map(|c| c.len_utf16() as u32).sum()
}

/// The highlighting of `text`, the file at `path`: `None` for a language
/// without a grammar here, or a file too large.
pub fn code(path: &str, text: &str) -> Option<Ranges> {
    if text.len() > LIMIT {
        return None;
    }
    let language = language_for(path, text)?;
    HIGHLIGHTER.with_borrow_mut(|highlighter| {
        let events = highlighter
            .highlight(&language.config, text.as_bytes(), None, None, |name| {
                by_name(name).map(|l| &l.config)
            })
            .ok()?;
        let mut ranges = Ranges::new(CODE);
        let mut stack: Vec<usize> = Vec::new();
        // Sources come in order, so offsets are counted once, as they go.
        let (mut byte, mut unit) = (0, 0u32);
        for event in events {
            match event.ok()? {
                HighlightEvent::HighlightStart(h) => stack.push(CAPTURES[h.0].1),
                HighlightEvent::HighlightEnd => {
                    stack.pop();
                }
                HighlightEvent::Source { start, end } => {
                    unit += utf16_len(&text[byte..start]);
                    let len = utf16_len(&text[start..end]);
                    if let Some(&class) = stack.last() {
                        ranges.push(class, unit, unit + len);
                    }
                    (byte, unit) = (end, unit + len);
                }
            }
        }
        Some(ranges)
    })
}

/// The highlighting of a file's changed lines, shown a row each: offsets
/// run over the lines' text as if joined by newlines. Removed lines are
/// highlighted as the old side reads, in its context, and the rest as the
/// new side reads; of a removed line and the added line paired with it,
/// the part that differs is marked too.
pub fn rows(path: &str, rows: &[Row]) -> Ranges {
    let lines: Vec<&Row> = rows.iter().filter(|r| r.kind.is_line()).collect();
    let mut starts = Vec::with_capacity(lines.len());
    let mut at = 0;
    for line in &lines {
        starts.push(at);
        at += utf16_len(&line.text) + 1;
    }
    let mut runs: Vec<Vec<(u32, u32)>> = vec![Vec::new(); ROWS.len()];
    for side in [Kind::Del, Kind::Add] {
        let other = if side == Kind::Del {
            Kind::Add
        } else {
            Kind::Del
        };
        let picked: Vec<usize> = (0..lines.len())
            .filter(|&i| lines[i].kind != other)
            .collect();
        let text: String = picked
            .iter()
            .map(|&i| format!("{}\n", lines[i].text))
            .collect();
        let Some(ranges) = code(path, &text) else {
            break;
        };
        let mut side_starts = Vec::with_capacity(picked.len());
        let mut at = 0;
        for &i in &picked {
            side_starts.push(at);
            at += utf16_len(&lines[i].text) + 1;
        }
        for (class, list) in ranges.runs.iter().enumerate() {
            for &(s, e) in list {
                let mut j = side_starts.partition_point(|&x| x <= s) - 1;
                while j < picked.len() && side_starts[j] < e {
                    let i = picked[j];
                    let (from, to) = (
                        s.max(side_starts[j]),
                        e.min(side_starts[j] + utf16_len(&lines[i].text)),
                    );
                    let shown = side == Kind::Add || lines[i].kind == Kind::Del;
                    if shown && from < to {
                        let base = starts[i] - side_starts[j];
                        runs[class].push((base + from, base + to));
                    }
                    j += 1;
                }
            }
        }
    }
    // Each run of removed lines, and the added lines after it.
    let mut i = 0;
    while i < lines.len() {
        let dels = lines[i..]
            .iter()
            .take_while(|l| l.kind == Kind::Del)
            .count();
        let adds = lines[i + dels..]
            .iter()
            .take_while(|l| l.kind == Kind::Add)
            .count();
        for k in 0..dels.min(adds) {
            let (d, a) = (i + k, i + dels + k);
            if let Some((dr, ar)) = changed(&lines[d].text, &lines[a].text) {
                runs[ROWS.len() - 1].push((starts[d] + dr.0, starts[d] + dr.1));
                runs[ROWS.len() - 2].push((starts[a] + ar.0, starts[a] + ar.1));
            }
        }
        i += (dels + adds).max(1);
    }
    let mut out = Ranges::new(ROWS);
    for (class, mut list) in runs.into_iter().enumerate() {
        list.sort_unstable();
        for (s, e) in list {
            out.push(class, s, e);
        }
    }
    out
}

/// Where `old` and `new` differ, in each, between what they start and end
/// with alike: `None` when that is less than half the shorter, as then the
/// line was rewritten rather than edited.
fn changed(old: &str, new: &str) -> Option<((u32, u32), (u32, u32))> {
    let (o, n): (Vec<char>, Vec<char>) = (old.chars().collect(), new.chars().collect());
    let prefix = o.iter().zip(&n).take_while(|(a, b)| a == b).count();
    let room = o.len().min(n.len()) - prefix;
    let suffix = o
        .iter()
        .rev()
        .zip(n.iter().rev())
        .take(room)
        .take_while(|(a, b)| a == b)
        .count();
    if (prefix + suffix) * 2 < o.len().min(n.len()) || prefix + suffix == 0 {
        return None;
    }
    let units = |c: &[char]| c.iter().map(|c| c.len_utf16() as u32).sum::<u32>();
    let span = |c: &[char]| (units(&c[..prefix]), units(&c[..c.len() - suffix]));
    Some((span(&o), span(&n)))
}

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

    /// The text each run of `name` covers.
    fn spans(r: &Ranges, text: &str, name: &str) -> Vec<String> {
        let units: Vec<u16> = text.encode_utf16().collect();
        let class = r.names.iter().position(|n| *n == name).unwrap();
        r.runs[class]
            .iter()
            .map(|(s, e)| String::from_utf16(&units[*s as usize..*e as usize]).unwrap())
            .collect()
    }

    #[test]
    fn rust() {
        let text = "// é😀 hi\nfn main() { let s: u8 = \"x\"; 42 }\n";
        let r = code("src/main.rs", text).unwrap();
        assert_eq!(spans(&r, text, "comment"), ["// é😀 hi"]);
        assert_eq!(spans(&r, text, "string"), ["\"x\""]);
        assert_eq!(spans(&r, text, "number"), ["42"]);
        assert_eq!(spans(&r, text, "function"), ["main"]);
        assert_eq!(spans(&r, text, "type"), ["u8"]);
        assert!(spans(&r, text, "keyword").contains(&"fn".to_string()));
        let encoded = r.encode();
        assert!(encoded.starts_with("comment:0,9;string:"), "{encoded}");
    }

    #[test]
    fn every_language_parses() {
        let samples = [
            ("a.js", "const f = () => <b x=\"1\">hi</b>; // c\n"),
            ("a.ts", "let x: number = 1;\n"),
            ("a.tsx", "const e = <div class=\"x\" />;\n"),
            ("a.py", "def f():\n    return 'x'\n"),
            ("a.go", "package main\nfunc main() {}\n"),
            ("a.c", "int main(void) { return 0; }\n"),
            ("a.cpp", "class A {};\n"),
            ("A.java", "class A { void f() {} }\n"),
            ("Gemfile", "gem 'x'\n"),
            ("a.json", "{\"key\": 1}\n"),
            ("Cargo.lock", "[[package]]\nname = \"x\"\n"),
            ("a.yml", "a: 'b'\n"),
            ("a.css", "a { color: red; }\n"),
            ("a.lua", "local x = 'y'\n"),
            ("a.zig", "const x = 1;\n"),
            ("run", "#!/usr/bin/env bash\necho 'hi'\n"),
            ("tool", "#!/usr/bin/python3.12\nx = 'y'\n"),
            ("cli", "#!/usr/bin/env node\nlet a = 'b'\n"),
            ("task", "#!/usr/bin/ruby\nputs 'x'\n"),
        ];
        for (path, text) in samples {
            let r = code(path, text).expect(path);
            assert!(!r.encode().is_empty(), "{path}");
        }
        let json = "{\"key\": 1}\n";
        assert_eq!(
            spans(&code("a.json", json).unwrap(), json, "attr"),
            ["\"key\""]
        );
    }

    #[test]
    fn html_carries_its_scripts_and_styles() {
        let html =
            "<a href=\"x\">y</a><script>let a = 'q';</script><style>b { color: red }</style>\n";
        let r = code("index.html", html).unwrap();
        assert_eq!(spans(&r, html, "attr"), ["href"]);
        assert!(
            spans(&r, html, "string").contains(&"'q'".to_string()),
            "{:?}",
            r
        );
        assert!(spans(&r, html, "tag").contains(&"script".to_string()));
    }

    #[test]
    fn what_it_leaves_plain() {
        assert_eq!(code("notes.unknown-kind", "plain\n"), None);
        assert_eq!(code("script", "#!/bin/perl\n"), None);
        assert_eq!(code("big.rs", &"x".repeat(LIMIT + 1)), None);
        let mut empty = Ranges::new(CODE);
        empty.push(0, 4, 4);
        assert_eq!(empty.encode(), "");
    }

    #[test]
    fn diff_rows() {
        let row = |kind, text: &str| Row {
            kind,
            old: None,
            new: None,
            text: text.into(),
        };
        let rows = [
            row(Kind::Hunk, "@@ -1,3 +1,3 @@"),
            row(Kind::Context, "/* a"),
            row(Kind::Del, "b */ let x = 1;"),
            row(Kind::Add, "b */ let x = 2;"),
            row(Kind::Add, "let é = \"q\";"),
            row(Kind::Note, "\\ No newline at end of file"),
            row(Kind::Del, "gone"),
            row(Kind::Add, "here"),
        ];
        let r = rows_named(&rows);
        let shown = "/* a\nb */ let x = 1;\nb */ let x = 2;\nlet é = \"q\";\ngone\nhere\n";
        // The comment spans two rows, and is cut at the row between.
        assert_eq!(spans(&r, shown, "comment"), ["/* a", "b */", "b */"]);
        assert_eq!(spans(&r, shown, "number"), ["1", "2"]);
        assert_eq!(spans(&r, shown, "string"), ["\"q\""]);
        assert_eq!(spans(&r, shown, "xdel"), ["1"]);
        assert_eq!(spans(&r, shown, "xins"), ["2"]);
        assert_eq!(
            super::rows("notes.unknown-kind", &rows).encode(),
            "xins:34,1;xdel:18,1"
        );
        assert_eq!(changed("abc", "xyz"), None);
        assert_eq!(changed("gone", "here"), None);
        assert_eq!(changed("aé", "aéb"), Some(((2, 2), (2, 3))));
        assert_eq!(changed("aa", "a"), Some(((1, 2), (1, 1))));
    }

    fn rows_named(rows: &[Row]) -> Ranges {
        super::rows("a.rs", rows)
    }
}