mod shell;
mod swift;
mod tree_sitter;

use once_cell::sync::Lazy;
use pulldown_cmark::escape::escape_html;
use tree_sitter::TreeSitterHighlighter;
use tree_sitter_highlight::HighlightConfiguration;

use self::{shell::ShellHighlighter, swift::SwiftHighlighter};

pub trait Highlighter {
    fn highlight(&self, code: &str) -> String;
}

pub fn highlight(code: &str, label: &str) -> String {
    match lang_highlighter(label) {
        None => {
            let mut buf = String::new();
            escape_html(&mut buf, &code).unwrap();
            buf
        }
        Some(highlighter) => highlighter.highlight(&code),
    }
}

macro_rules! ts_config {
    ($name:ident, $module:ident, $query:ident) => {
        static $name: Lazy<HighlightConfiguration> = Lazy::new(|| {
            let mut config =
                HighlightConfiguration::new($module::language(), $module::$query, "", "").unwrap();
            config.configure(tree_sitter::HIGHLIGHT_NAMES);
            config
        });
    };
    ($name:ident, $module:ident, $query:literal) => {
        static $name: Lazy<HighlightConfiguration> = Lazy::new(|| {
            let mut config =
                HighlightConfiguration::new($module::language(), include_str!($query), "", "")
                    .unwrap();
            config.configure(tree_sitter::HIGHLIGHT_NAMES);
            config
        });
    };
}

ts_config!(BASH_CONFIG, tree_sitter_bash, HIGHLIGHTS_QUERY);
ts_config!(C_CONFIG, tree_sitter_c, HIGHLIGHT_QUERY);
ts_config!(CSS_CONFIG, tree_sitter_css, "highlight/css-highlight.scm");
ts_config!(ELIXIR_CONFIG, tree_sitter_elixir, HIGHLIGHTS_QUERY);
ts_config!(HTML_CONFIG, tree_sitter_html, HIGHLIGHTS_QUERY);
ts_config!(JAVA_CONFIG, tree_sitter_java, HIGHLIGHT_QUERY);
ts_config!(JS_CONFIG, tree_sitter_javascript, HIGHLIGHT_QUERY);
ts_config!(JSON_CONFIG, tree_sitter_json, HIGHLIGHT_QUERY);
ts_config!(
    OBJC_CONFIG,
    tree_sitter_objc,
    "highlight/objc-highlight.scm"
);
ts_config!(RUST_CONFIG, tree_sitter_rust, HIGHLIGHT_QUERY);

fn lang_highlighter(name: &str) -> Option<Box<dyn Highlighter>> {
    match name {
        "" | "plaintext" | "plain" | "txt" => None,
        "c" => Some(Box::new(TreeSitterHighlighter::new(&*C_CONFIG))),
        "css" => Some(Box::new(TreeSitterHighlighter::new(&*CSS_CONFIG))),
        "elixir" => Some(Box::new(TreeSitterHighlighter::new(&*ELIXIR_CONFIG))),
        "html" => Some(Box::new(TreeSitterHighlighter::new(&*HTML_CONFIG))),
        "java" => Some(Box::new(TreeSitterHighlighter::new(&*JAVA_CONFIG))),
        "js" | "javascript" => Some(Box::new(TreeSitterHighlighter::new(&*JS_CONFIG))),
        "json" => Some(Box::new(TreeSitterHighlighter::new(&*JSON_CONFIG))),
        "objc" | "objective-c" | "objective-c++" => {
            Some(Box::new(TreeSitterHighlighter::new(&*OBJC_CONFIG)))
        }
        "rust" => Some(Box::new(TreeSitterHighlighter::new(&RUST_CONFIG))),
        "sh" | "bash" => Some(Box::new(TreeSitterHighlighter::new(&*BASH_CONFIG))),
        "shell" => Some(Box::new(ShellHighlighter)),
        "swift" => Some(Box::new(SwiftHighlighter)),
        _ => panic!("unsupported language {}", name),
    }
}

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

    #[test]
    fn test_highlight_plain() {
        assert_eq!(highlight("foo(bar, 123)".into(), ""), "foo(bar, 123)");
    }

    #[test]
    fn test_highlight_js() {
        assert_eq!(
            highlight("foo(bar, 123)".into(), "js"),
            r#"<span class="hl-fn">foo</span>(<span class="hl-var">bar</span>, <span class="hl-num">123</span>)"#
        );
    }
}
