use iced::advanced::text;
use iced::advanced::text::highlighter;
use iced::{Color, Font};
use std::ops::Range;

const KW: &[&str] = &[
    "let", "fn", "if", "else", "for", "in", "while", "return", "true", "false",
    "cast", "plot", "sch", "map",
];

const BUILTINS: &[&str] = &[
    "sin", "cos", "tan", "asin", "acos", "atan", "atan2",
    "sinh", "cosh", "tanh", "asinh", "acosh", "atanh",
    "arcsin", "arccos", "arctan", "arcsinh", "arccosh", "arctanh", "arcos", "arcosh",
    "sqrt", "exp", "ln", "log", "abs", "hypot", "min", "max",
    "length", "mag", "mix", "clip", "clamp", "smoothstep", "quantize",
    "saw", "tri", "square",
    "lpf", "hpf", "bpf", "am", "fm", "dft",
    "envelope", "hilbert", "phase",
    "sphere", "box", "cylinder", "ngon",
    "translate", "mov", "move",
    "rotate_x", "rotate_y", "rotate_z", "rx", "ry", "rz",
    "scale", "mirror_x", "mirror_y", "mirror_z", "mx", "my", "mz",
    "union", "intersect", "diff", "subtract",
];

const TYPES: &[&str] = &[
    "f64", "f32", "i32", "u32", "bool", "sdf", "vec2", "vec3", "vec4",
    "Obj", "obj",
];

const CONSTS: &[&str] = &["pi", "PI", "e", "E", "x", "y", "z", "reg"];

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TokenKind {
    Keyword,
    Builtin,
    Constant,
    TypeName,
    Number,
    Operator,
    Paren,
    Comment,
    Plain,
}

#[derive(Debug, Clone, PartialEq)]
pub struct CordHighlighterSettings;

pub struct CordHighlighter {
    current_line: usize,
    in_block_comment: bool,
}

impl text::Highlighter for CordHighlighter {
    type Settings = CordHighlighterSettings;
    type Highlight = TokenKind;
    type Iterator<'a> = std::vec::IntoIter<(Range<usize>, TokenKind)>;

    fn new(_settings: &Self::Settings) -> Self {
        Self { current_line: 0, in_block_comment: false }
    }

    fn update(&mut self, _new_settings: &Self::Settings) {}

    fn change_line(&mut self, line: usize) {
        self.current_line = line;
        if line == 0 {
            self.in_block_comment = false;
        }
    }

    fn highlight_line(&mut self, line: &str) -> Self::Iterator<'_> {
        self.current_line += 1;
        let spans = lex_line_with_state(line, &mut self.in_block_comment);
        spans.into_iter()
    }

    fn current_line(&self) -> usize {
        self.current_line
    }
}

pub fn format_token(kind: &TokenKind, _theme: &iced::Theme) -> highlighter::Format<Font> {
    let color = match kind {
        TokenKind::Keyword  => Color::from_rgb(0.55, 0.75, 1.0),
        TokenKind::Builtin  => Color::from_rgb(0.6, 0.85, 0.75),
        TokenKind::Constant => Color::from_rgb(0.85, 0.7, 1.0),
        TokenKind::TypeName => Color::from_rgb(0.45, 0.80, 0.95),
        TokenKind::Number   => Color::from_rgb(0.95, 0.75, 0.45),
        TokenKind::Operator => Color::from_rgb(0.85, 0.85, 0.85),
        TokenKind::Paren    => Color::from_rgb(0.65, 0.65, 0.65),
        TokenKind::Comment  => Color::from_rgb(0.45, 0.50, 0.45),
        TokenKind::Plain    => Color::from_rgb(0.90, 0.90, 0.90),
    };
    highlighter::Format {
        color: Some(color),
        font: None,
    }
}

fn lex_line_with_state(line: &str, in_block_comment: &mut bool) -> Vec<(Range<usize>, TokenKind)> {
    let mut spans = Vec::new();
    let bytes = line.as_bytes();
    let len = bytes.len();
    let mut i = 0;

    while i < len {
        // Inside block comment: scan for */
        if *in_block_comment {
            let start = i;
            while i < len {
                if bytes[i] == b'*' && i + 1 < len && bytes[i + 1] == b'/' {
                    i += 2;
                    *in_block_comment = false;
                    break;
                }
                i += 1;
            }
            spans.push((start..i, TokenKind::Comment));
            continue;
        }

        let b = bytes[i];

        // Line comment (// or /=)
        if b == b'/' && i + 1 < len && (bytes[i + 1] == b'/' || bytes[i + 1] == b'=') {
            spans.push((i..len, TokenKind::Comment));
            break;
        }

        // Block comment start
        if b == b'/' && i + 1 < len && bytes[i + 1] == b'*' {
            let start = i;
            i += 2;
            while i < len {
                if bytes[i] == b'*' && i + 1 < len && bytes[i + 1] == b'/' {
                    i += 2;
                    break;
                }
                i += 1;
            }
            if i >= len && !(len >= 2 && bytes[len - 2] == b'*' && bytes[len - 1] == b'/') {
                *in_block_comment = true;
            }
            spans.push((start..i, TokenKind::Comment));
            continue;
        }

        // Whitespace
        if b == b' ' || b == b'\t' {
            let start = i;
            while i < len && (bytes[i] == b' ' || bytes[i] == b'\t') {
                i += 1;
            }
            spans.push((start..i, TokenKind::Plain));
            continue;
        }

        // Number
        if b.is_ascii_digit() || (b == b'.' && i + 1 < len && bytes[i + 1].is_ascii_digit()) {
            let start = i;
            while i < len && (bytes[i].is_ascii_digit() || bytes[i] == b'.') {
                i += 1;
            }
            spans.push((start..i, TokenKind::Number));
            continue;
        }

        // Identifier
        if b.is_ascii_alphabetic() || b == b'_' {
            let start = i;
            while i < len && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') {
                i += 1;
            }
            let word = &line[start..i];
            let kind = if KW.contains(&word) {
                TokenKind::Keyword
            } else if TYPES.contains(&word) {
                TokenKind::TypeName
            } else if BUILTINS.contains(&word) {
                TokenKind::Builtin
            } else if CONSTS.contains(&word) {
                TokenKind::Constant
            } else {
                TokenKind::Plain
            };
            spans.push((start..i, kind));
            continue;
        }

        // Operators
        if b"+-*/^=;:".contains(&b) {
            spans.push((i..i + 1, TokenKind::Operator));
            i += 1;
            continue;
        }

        // Parens
        if b"()[]{}".contains(&b) {
            spans.push((i..i + 1, TokenKind::Paren));
            i += 1;
            continue;
        }

        // Comma
        if b == b',' {
            spans.push((i..i + 1, TokenKind::Plain));
            i += 1;
            continue;
        }

        spans.push((i..i + 1, TokenKind::Plain));
        i += 1;
    }

    spans
}
