use std::collections::HashMap;

use crate::{
    BString, ansi, bstr, cursor,
    parse::{Highlight, HighlightKind},
};

pub struct Highlighter {
    pub enabled: bool,
    colors: HashMap<HighlightKind, BString>,
}

#[derive(Debug)]
pub enum SetColorError {
    NoSuchKeyword,
}

const CLEAR_TO_END: &bstr = b"\x1b[K";

impl Highlighter {
    pub fn new() -> Self {
        let mut this = Self {
            enabled: false,
            colors: HashMap::new(),
        };
        let mut sc = |a: &str, b: &str| this.set_color(a.as_bytes(), b.as_bytes()).unwrap();
        use ansi::colors::*;
        sc("keywords", GREEN_FG);
        sc("braces", CYAN_FG);
        sc("string", MAGENTA_FG);
        sc("var", CYAN_FG);
        sc("escape", RED_FG);
        this
    }

    pub fn set_color(&mut self, ident: &bstr, color: &bstr) -> Result<(), SetColorError> {
        let kinds = HighlightKind::from_identifier(ident);
        for kind in kinds.iter() {
            self.colors.insert(*kind, color.to_vec());
        }
        if kinds.is_empty() {
            Err(SetColorError::NoSuchKeyword)
        } else {
            Ok(())
        }
    }

    pub fn color(&self, h: HighlightKind) -> &[u8] {
        self.colors
            .get(&h)
            .map(|c| c.as_ref())
            .unwrap_or(ansi::colors::RESET.as_bytes())
    }

    pub fn pretty_print(
        &self,
        bytes: &[u8],
        cursor_offset: usize,
        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);

        let mut saved_pos = false;

        stdout.write_all(CLEAR_TO_END)?;

        for (i, x) in bytes.iter().cloned().enumerate() {
            if i == cursor_offset {
                cursor::f_save(stdout)?;
                saved_pos = true;
            }

            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;
                }
            }

            if x == b'\n' {
                stdout.write_all(b"\r\n")?;
                stdout.write_all(CLEAR_TO_END)?;
            } else {
                stdout.write_all(&[x])?;
            }
        }

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

        if saved_pos {
            cursor::f_restore(stdout)?;
        }

        Ok(())
    }
}