mod highlights;
mod parsers;

use core::str;
use std::{
    cell::RefCell,
    collections::HashMap,
    env::{
        args,
        current_dir,
    },
    fs::{
        create_dir,
        read,
    },
    process::Command,
};

use dirs::cache_dir;
use highlights::HIGHLIGHTS;
use libloading::Library;
use parsers::{
    ALIASES,
    PARSERS,
};
use tree_sitter::Language;
use tree_sitter_highlight::{
    HighlightConfiguration, Highlighter, HtmlRenderer
};
use url::Url;

fn get_config<'a>(
    lang: &'a str,
    configs: &'a RefCell<HashMap<&'a str, &'static HighlightConfiguration>>,
) -> Option<&'static HighlightConfiguration> {
    let lang = ALIASES.get(lang).unwrap_or(&lang);
    configs.borrow_mut().entry(lang).or_insert_with(|| {
        let install_info = PARSERS.get(lang).unwrap();
        let mut cache_dir = cache_dir().unwrap();
        cache_dir.push("yah");

        let mut query_dir = current_dir().unwrap();
        query_dir.push("nvim-treesitter");
        query_dir.push("queries");
        query_dir.push(lang);

        let mut parsers_dir = current_dir().unwrap();
        parsers_dir.push(".parsers");
        if !parsers_dir.exists() {
            create_dir(&parsers_dir).unwrap();
        }
        if !cache_dir.exists() {
            create_dir(&cache_dir).unwrap();
        }

        let parser_path = parsers_dir.join(format!("{lang}.so"));
        let repo_path = cache_dir.join(
            Url::parse(install_info.url)
                .unwrap()
                .path_segments()
                .unwrap()
                .last()
                .unwrap(),
        );

        if !parser_path.exists() {
            if !repo_path.exists() {
                Command::new("git")
                    .arg("clone")
                    .arg(install_info.url)
                    .current_dir(&cache_dir)
                    .spawn()
                    .unwrap()
                    .wait()
                    .unwrap();
            }

            if install_info.requires_generate_from_grammar {
                if install_info.generate_requires_npm {
                    Command::new("npm")
                        .arg("install")
                        .current_dir(&repo_path)
                        .spawn()
                        .unwrap()
                        .wait()
                        .unwrap();
                }

                Command::new("tree-sitter")
                    .arg("generate")
                    .arg("--no-bindings")
                    .current_dir(&repo_path)
                    .spawn()
                    .unwrap()
                    .wait()
                    .unwrap();
            }

            Command::new("tree-sitter")
                .arg("build")
                .arg("-o")
                .arg(&parser_path)
                .arg(
                    install_info
                        .location
                        .map_or(repo_path.clone(), |location| repo_path.join(location)),
                )
                .spawn()
                .unwrap()
                .wait()
                .unwrap();
        }

        let library = unsafe {
            Library::new(parser_path)
                .unwrap()
        };
        let parser = unsafe {
            library.get::<unsafe extern "C" fn() -> Language>(format!("tree_sitter_{lang}").as_bytes())
                .unwrap()
                ()
        };
        std::mem::forget(library); // this causes the dylib to not be unloaded with dlclose after this
        // scope is dropped, and thus extending its lifetime to static

        let mut config = HighlightConfiguration::new(
            parser,
            *lang,
            read(query_dir.join("highlights.scm"))
                .map(|v| String::from_utf8(v).unwrap())
                .unwrap_or(String::from(""))
                .as_str(),
            read(query_dir.join("injections.scm"))
                .map(|v| String::from_utf8(v).unwrap())
                .unwrap_or(String::from(""))
                .as_str(),
            read(query_dir.join("locals.scm"))
                .map(|v| String::from_utf8(v).unwrap())
                .unwrap_or(String::from(""))
                .as_str(),
        )
        .unwrap();
        config.configure(HIGHLIGHTS);
        Box::leak(Box::new(config))
    });
    configs.borrow_mut().get(lang).map(|v| *v)
}

fn main() {
    if let [_, lang, code] = args().collect::<Vec<String>>().as_slice() {
        let mut highlighter = Highlighter::new();
        let configs = RefCell::new(HashMap::new());

        let config = get_config(lang, &configs).unwrap();

        let events = highlighter
            .highlight(&config, code.as_bytes(), None, |v| {
                let lang = String::from(v).leak();
                get_config(lang, &configs)
            })
            .unwrap();

        let mut renderer = HtmlRenderer::new();
        renderer.render(events, code.as_bytes(), &|highlight, output| {
            output.extend(b"class=\"");
            let mut parts = HIGHLIGHTS[highlight.0].split('.').peekable();
            while let Some(part) = parts.next() {
                output.extend(part.as_bytes());
                if parts.peek().is_some() {
                    output.extend(b" ");
                }
            }
            output.extend(b"\"");
        }).unwrap();
        print!("{}", String::from_utf8(renderer.html).unwrap())
    } else {
        panic!("Need <lang> <code> as arguments")
    }
}