use inkjet::constants::HIGHLIGHT_NAMES;
use inkjet::formatter::{Formatter, Theme};
use inkjet::tree_sitter_highlight::HighlightEvent;
use inkjet::{Highlighter, Language, Result as InkjetResult};
use std::io::{self, Write};

use std::cell::RefCell;
use std::fmt::Write as FmtWrite;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};

use crate::utils::theme::create_theme;

pub fn highlight_code<W: Write>(code: &str, lang: &str, writer: &mut W) -> io::Result<()> {
    let mut highlighter = Highlighter::new();
    let language = Language::from_token(lang).unwrap_or(Language::Plaintext);
    let theme = create_theme();
    let formatter = TerminalFormatter::new(theme);

    for line in code.lines() {
        highlighter
            .highlight_to_writer(language, &formatter, line, writer)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
    }

    Ok(())
}

struct TerminalFormatter {
    theme: Theme,
    stdout: RefCell<StandardStream>,
}

impl TerminalFormatter {
    fn new(theme: Theme) -> Self {
        Self {
            theme,
            stdout: RefCell::new(StandardStream::stdout(ColorChoice::Always)),
        }
    }

    fn color_from_hex(&self, hex: &str) -> Color {
        let rgb = color_from_hex(hex).unwrap_or((255, 255, 255));
        Color::Rgb(rgb.0, rgb.1, rgb.2)
    }
}

impl Formatter for TerminalFormatter {
    fn write<W>(&self, source: &str, _writer: &mut W, event: HighlightEvent) -> InkjetResult<()>
    where
        W: FmtWrite,
    {
        match event {
            HighlightEvent::Source { start, end } => {
                let text = &source[start..end];
                let mut stdout = self.stdout.borrow_mut();
                stdout.write_all(text.as_bytes())?;
                stdout.flush()?;
            }
            HighlightEvent::HighlightStart(highlight) => {
                let style_name = HIGHLIGHT_NAMES[highlight.0];
                let style = self.theme.get_style(style_name);
                let color = self.color_from_hex(&style.primary_color);
                self.stdout
                    .borrow_mut()
                    .set_color(ColorSpec::new().set_fg(Some(color)))?;
            }
            HighlightEvent::HighlightEnd => {
                self.stdout.borrow_mut().reset()?;
            }
        }
        Ok(())
    }
}

fn color_from_hex(hex: &str) -> Option<(u8, u8, u8)> {
    if hex.len() != 7 || !hex.starts_with('#') {
        return None;
    }

    let r = u8::from_str_radix(&hex[1..3], 16).ok()?;
    let g = u8::from_str_radix(&hex[3..5], 16).ok()?;
    let b = u8::from_str_radix(&hex[5..7], 16).ok()?;

    Some((r, g, b))
}