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<()> {
        #[derive(PartialEq, Eq, Debug, Clone)]
        struct ColorBoundary {
            loc: usize,
            is_end: bool,
            kind: HighlightKind,
        }

        impl PartialOrd for ColorBoundary {
            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
                Some(self.cmp(other))
            }
        }

        impl Ord for ColorBoundary {
            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
                match self.loc.cmp(&other.loc) {
                    std::cmp::Ordering::Equal => (),
                    ord => return ord,
                }
                match self.kind.cmp(&other.kind) {
                    std::cmp::Ordering::Equal => (),
                    ord => return ord,
                }
                self.is_end.cmp(&other.is_end)
            }
        }

        let mut coloring: Vec<_> = colors
            .into_iter()
            .flat_map(|hi| {
                [
                    ColorBoundary {
                        loc: hi.span.start as usize,
                        is_end: false,
                        kind: hi.kind,
                    },
                    ColorBoundary {
                        loc: hi.span.end as usize,
                        is_end: true,
                        kind: hi.kind,
                    },
                ]
            })
            .collect();
        coloring.sort();
        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(color_boundary) = coloring.first().cloned()
                && color_boundary.loc <= i
            {
                coloring = &coloring[1..];
                if color_boundary.is_end {
                    color_stack.pop();
                } else {
                    color_stack.push(color_boundary.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(())
    }
}