use html_escape::encode_text;
use tree_sitter_highlight::{HighlightEvent, Highlighter};

mod languages;

pub fn highlight(source: &str) -> String {
    let source = source.as_bytes();

    let mut highlighter = Highlighter::new();
    let (highlight_config, capture_names) =
        languages::get_highlight_config("rust").expect("Could not get Rust language config");

    let items = highlighter
        .highlight(&highlight_config, source, None, |lang| {
            languages::get_highlight_config(lang).map(|x| /* ugh */ &*Box::leak(Box::new(x.0)))
        })
        .unwrap();

    let mut highlight_attrs: Vec<&str> = Vec::new();
    let mut out = String::new();

    // Collapse adjacent identical attribute sets
    let mut last_attrs: u64 = 0;
    let mut span_is_open = false;

    for item in items {
        match item.unwrap() {
            HighlightEvent::Source { start, end } => {
                let source_section_bytes = &source[start..end];
                let source_section = String::from_utf8_lossy(source_section_bytes);

                let attr_hash = {
                    use std::{
                        collections::hash_map::DefaultHasher,
                        hash::{Hash, Hasher},
                    };

                    let mut hasher = DefaultHasher::new();
                    highlight_attrs.hash(&mut hasher);
                    hasher.finish()
                };

                if last_attrs != attr_hash {
                    if span_is_open {
                        out.push_str("</span>");
                        span_is_open = false;
                    }

                    if !highlight_attrs.is_empty() {
                        out.push_str("<span class=\"");
                        for (i, attr) in highlight_attrs.iter().enumerate() {
                            if i != 0 {
                                out.push(' ');
                            }
                            out.push_str(&attr.replace('.', " "));
                        }
                        out.push_str("\">");
                        span_is_open = true;
                        last_attrs = attr_hash;
                    }
                }

                out.push_str(&encode_text(&source_section));
            }
            HighlightEvent::HighlightStart(highlight) => {
                let capture_name = &capture_names[highlight.0];
                highlight_attrs.push(capture_name);
            }
            HighlightEvent::HighlightEnd => {
                highlight_attrs.pop();
            }
        }
    }

    if span_is_open {
        out.push_str("</span>");
    }

    out
}

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

    #[test]
    fn highlight_rust() {
        println!(
            "{}",
            highlight(
                r#"fn main() {
    println!("Hello, world!");
}"#
            )
        )
    }
}