use crate::parse::{Highlight, HighlightKind, Keyword};

pub struct Highlighter {
    pub enabled: bool,
}

impl Highlighter {
    pub fn new() -> Self {
        Self { enabled: true }
    }

    pub fn color(&self, h: HighlightKind) -> &[u8] {
        // TODO: configurable
        const GREEN: &[u8] = b"\x1b[32m";
        const BLUE: &[u8] = b"\x1b[36m";
        const MAGENTA: &[u8] = b"\x1b[95m";
        const COLOR_RESET: &[u8] = b"\x1b[0m";
        match h {
            HighlightKind::Keyword(
                Keyword::If | Keyword::Elif | Keyword::Else | Keyword::While,
            ) => GREEN,
            HighlightKind::Keyword(Keyword::OpenBrace | Keyword::CloseBrace) => BLUE,
            HighlightKind::String => MAGENTA,
            HighlightKind::None => COLOR_RESET,
        }
    }

    pub fn pretty_print(
        &self,
        bytes: &[u8],
        colors: Vec<Highlight>,
        stdout: &mut dyn std::io::Write,
    ) -> std::io::Result<()> {
        let mut coloring: Vec<_> = colors
            .into_iter()
            .flat_map(|hi| {
                [
                    (hi.span.start as usize, false, hi.kind),
                    (hi.span.end as usize, true, hi.kind),
                ]
            })
            .collect();
        coloring.sort_by_key(|x| (x.0, x.1));
        let mut coloring = &coloring[..];
        let mut color_stack = Vec::new();

        let mut current_color = self.color(HighlightKind::None);

        for (i, x) in bytes.iter().cloned().enumerate() {
            while let Some((k, is_end, kind)) = coloring.first().cloned()
                && k == i
            {
                coloring = &coloring[1..];
                if is_end {
                    color_stack.pop();
                } else {
                    color_stack.push(kind);
                }

                let new_color =
                    self.color(color_stack.last().cloned().unwrap_or(HighlightKind::None));

                if current_color != new_color {
                    stdout.write_all(new_color)?;
                    current_color = new_color;
                }
            }

            stdout.write_all(&[x])?;
        }

        stdout.write_all(self.color(HighlightKind::None))?;

        Ok(())
    }
}