//! Syntax highlighting derived from the same Chumsky parser and lexer.
//! Single source of truth: no separate TextMate grammar.

use std::collections::HashSet;

use crate::parse::ast::{Expr, FuncDef, Span, Spanned};
use crate::parse::lexer::TokenKind;
use crate::parse::{lex_with_offset, strip_line_comments_preserve_len};

// ---------------------------------------------------------------------------
// Line index and LSP semantic tokens
// ---------------------------------------------------------------------------

/// Build a line index: `line_starts[i]` is the byte offset of the start of line `i` (0-based).
pub fn line_starts(source: &str) -> Vec<usize> {
    let mut starts = vec![0];
    for (i, b) in source.bytes().enumerate() {
        if b == b'\n' {
            starts.push(i + 1);
        }
    }
    starts
}

/// Convert byte offset to (line, column) 0-based. Column is in UTF-8 bytes for LSP.
pub fn byte_offset_to_line_col(offset: usize, line_starts: &[usize]) -> (u32, u32) {
    let line = line_starts
        .partition_point(|&s| s <= offset)
        .saturating_sub(1);
    let line_start = line_starts.get(line).copied().unwrap_or(0);
    let col = (offset - line_start) as u32;
    (line as u32, col)
}

/// LSP semantic token type indices. Must match the legend returned by `semantic_tokens_legend()`.
pub const LSP_TOKEN_TYPE_KEYWORD: u32 = 0;
pub const LSP_TOKEN_TYPE_VARIABLE: u32 = 1;
pub const LSP_TOKEN_TYPE_FUNCTION: u32 = 2;
pub const LSP_TOKEN_TYPE_TYPE: u32 = 3;
pub const LSP_TOKEN_TYPE_NUMBER: u32 = 4;
pub const LSP_TOKEN_TYPE_CONSTANT: u32 = 5;
pub const LSP_TOKEN_TYPE_OPERATOR: u32 = 6;
pub const LSP_TOKEN_TYPE_PUNCTUATION: u32 = 7;
pub const LSP_TOKEN_TYPE_COMMENT: u32 = 8;

fn highlight_kind_to_lsp_type(kind: HighlightKind) -> u32 {
    match kind {
        HighlightKind::Keyword => LSP_TOKEN_TYPE_KEYWORD,
        HighlightKind::Variable => LSP_TOKEN_TYPE_VARIABLE,
        HighlightKind::Function => LSP_TOKEN_TYPE_FUNCTION,
        HighlightKind::Type => LSP_TOKEN_TYPE_TYPE,
        HighlightKind::Number => LSP_TOKEN_TYPE_NUMBER,
        HighlightKind::Constant => LSP_TOKEN_TYPE_CONSTANT,
        HighlightKind::Operator => LSP_TOKEN_TYPE_OPERATOR,
        HighlightKind::Punctuation => LSP_TOKEN_TYPE_PUNCTUATION,
        HighlightKind::Comment => LSP_TOKEN_TYPE_COMMENT,
    }
}

/// Token types and modifiers for LSP `SemanticTokensLegend`. Order of token_types
/// must match the constants `LSP_TOKEN_TYPE_*`.
pub fn semantic_tokens_legend() -> (Vec<String>, Vec<String>) {
    let token_types = vec![
        "keyword".into(),
        "variable".into(),
        "function".into(),
        "type".into(),
        "number".into(),
        "constant".into(),
        "operator".into(),
        "punctuation".into(),
        "comment".into(),
    ];
    let token_modifiers = vec![];
    (token_types, token_modifiers)
}

/// Convert `(Span, HighlightKind)` list to LSP semantic token array.
/// Each token is 5 u32s: [deltaLine, deltaStartChar, length, tokenType, tokenModifiers].
/// Delta is relative to the previous token's (line, startChar).
pub fn to_semantic_tokens(spans: &[(Span, HighlightKind)], source: &str) -> Vec<u32> {
    let line_starts = line_starts(source);
    let mut out = Vec::with_capacity(spans.len() * 5);
    let mut prev_line = 0u32;
    let mut prev_start_char = 0u32;
    for (span, kind) in spans {
        let (line, start_char) = byte_offset_to_line_col(span.start, &line_starts);
        let length = (span.end.saturating_sub(span.start)) as u32;
        let delta_line = line - prev_line;
        let delta_start = if line == prev_line {
            start_char - prev_start_char
        } else {
            start_char
        };
        out.push(delta_line);
        out.push(delta_start);
        out.push(length);
        out.push(highlight_kind_to_lsp_type(*kind));
        out.push(0); // no modifiers
        prev_line = line;
        prev_start_char = start_char;
    }
    out
}

/// Kind of highlight for a span, maps to LSP/theme categories.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum HighlightKind {
    Keyword,
    Variable,
    Function,
    Type,
    Number,
    Constant,
    Operator,
    Punctuation,
    Comment,
}

/// Scan source for line comments (`#` to newline) and emit (span, Comment).
fn scan_comments(source: &str) -> Vec<(Span, HighlightKind)> {
    let mut out = Vec::new();
    let bytes = source.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'#' {
            let start = i;
            i += 1;
            while i < bytes.len() && bytes[i] != b'\n' {
                i += 1;
            }
            out.push((Span::new(start, i), HighlightKind::Comment));
        } else {
            i += 1;
        }
    }
    out
}

/// Find header `|...|` and emit (span, Variable) for param names, (span, Type) for param types.
fn highlight_header(source: &str) -> Vec<(Span, HighlightKind)> {
    let start = match source.find('|') {
        Some(s) => s,
        None => return Vec::new(),
    };
    let after_start = start + 1;
    let end_rel = match source[after_start..].find('|') {
        Some(r) => r,
        None => return Vec::new(),
    };
    let end = after_start + end_rel;
    let params_src = &source[after_start..end];
    let mut out = Vec::new();
    let mut cursor = after_start;
    for chunk in params_src.split(',') {
        let trimmed = chunk.trim();
        let leading_ws = chunk.len() - chunk.trim_start().len();
        let chunk_start = cursor + leading_ws;
        let chunk_end = chunk_start + trimmed.len();
        if let Some((name, dtype_str)) = trimmed.split_once(':') {
            let name_trimmed = name.trim();
            let dtype_trimmed = dtype_str.trim();
            if !name_trimmed.is_empty() {
                out.push((
                    Span::new(chunk_start, chunk_start + name_trimmed.len()),
                    HighlightKind::Variable,
                ));
            }
            if !dtype_trimmed.is_empty() {
                out.push((
                    Span::new(chunk_end - dtype_trimmed.len(), chunk_end),
                    HighlightKind::Type,
                ));
            }
        }
        cursor += chunk.len() + 1;
    }
    out
}

/// Map TokenKind to HighlightKind for lexical (non-semantic) tokens.
fn token_kind_to_highlight(tok: &TokenKind) -> HighlightKind {
    match tok {
        TokenKind::If | TokenKind::Then | TokenKind::Else | TokenKind::Cast => {
            HighlightKind::Keyword
        }
        TokenKind::True | TokenKind::False | TokenKind::ConstE | TokenKind::ConstPi => {
            HighlightKind::Constant
        }
        TokenKind::Ident(_) => HighlightKind::Variable,
        TokenKind::DTypeName(_) => HighlightKind::Type,
        TokenKind::Int(_) | TokenKind::Float(_) | TokenKind::Complex(_) => HighlightKind::Number,
        TokenKind::Plus
        | TokenKind::Minus
        | TokenKind::Star
        | TokenKind::Slash
        | TokenKind::SlashSlash
        | TokenKind::Percent
        | TokenKind::Pow
        | TokenKind::AndAnd
        | TokenKind::OrOr
        | TokenKind::Bang
        | TokenKind::EqEq
        | TokenKind::Ne
        | TokenKind::Le
        | TokenKind::Ge
        | TokenKind::Lt
        | TokenKind::Gt
        | TokenKind::At => HighlightKind::Operator,
        TokenKind::LParen | TokenKind::RParen | TokenKind::Comma => HighlightKind::Punctuation,
    }
}

/// Collect (span, kind) from AST for semantic highlighting.
fn collect_semantic_from_expr(expr: &Spanned<Expr>, out: &mut Vec<(Span, HighlightKind)>) {
    match &expr.node {
        Expr::Var(_) => out.push((expr.span, HighlightKind::Variable)),
        Expr::LitBool(_)
        | Expr::LitInt(_, _)
        | Expr::LitFloat(_, _)
        | Expr::LitComplex(_, _, _) => out.push((expr.span, HighlightKind::Number)),
        Expr::ConstE | Expr::ConstPi => out.push((expr.span, HighlightKind::Constant)),
        Expr::Call(_, args) => {
            out.push((expr.span, HighlightKind::Function));
            for a in args {
                collect_semantic_from_expr(a, out);
            }
        }
        Expr::Cast(_, inner) => {
            out.push((expr.span, HighlightKind::Type));
            collect_semantic_from_expr(inner, out);
        }
        Expr::Unary(_, inner) => collect_semantic_from_expr(inner, out),
        Expr::Binary(_, lhs, rhs) => {
            collect_semantic_from_expr(lhs, out);
            collect_semantic_from_expr(rhs, out);
        }
        Expr::IfElse(cond, then_e, else_e) => {
            collect_semantic_from_expr(cond, out);
            collect_semantic_from_expr(then_e, out);
            collect_semantic_from_expr(else_e, out);
        }
        Expr::Tuple(items) => {
            for item in items {
                collect_semantic_from_expr(item, out);
            }
        }
        Expr::Let(_, value, body) => {
            collect_semantic_from_expr(value, out);
            collect_semantic_from_expr(body, out);
        }
    }
}

/// Merge semantic and lexical: for each token span, use semantic kind if that span was
/// produced by AST (exact match), else use lexical kind. Semantic takes precedence so
/// e.g. variable names get Variable not default Ident→Variable.
fn merge_semantic_lexical(
    semantic: Vec<(Span, HighlightKind)>,
    tokens: Vec<(TokenKind, Span)>,
) -> Vec<(Span, HighlightKind)> {
    let semantic_set: HashSet<(usize, usize)> =
        semantic.iter().map(|(s, _)| (s.start, s.end)).collect();
    let mut out = Vec::with_capacity(semantic.len() + tokens.len());
    for (span, kind) in semantic {
        out.push((span, kind));
    }
    for (tok, span) in tokens {
        if !semantic_set.contains(&(span.start, span.end)) {
            out.push((span, token_kind_to_highlight(&tok)));
        }
    }
    out.sort_by_key(|(s, _)| s.start);
    out
}

/// Main entry: derive highlights from source and optional parsed FuncDef.
/// When func is None (e.g. parse failed), only comments and lexical body/header are produced.
pub fn highlight(source: &str, func: Option<&FuncDef>) -> Vec<(Span, HighlightKind)> {
    let mut out = Vec::new();

    // Comments (raw source)
    out.extend(scan_comments(source));

    // Header param names / types
    out.extend(highlight_header(source));

    // Body: semantic from AST + lexical from token stream
    let cleaned = strip_line_comments_preserve_len(source);
    let body_start = match cleaned.find('|') {
        None => return out,
        Some(first) => match cleaned[first + 1..].find('|') {
            None => return out,
            Some(rel) => first + 1 + rel + 1,
        },
    };
    if body_start >= cleaned.len() {
        return out;
    }
    let body_slice = &cleaned[body_start..];

    let semantic: Vec<(Span, HighlightKind)> = if let Some(f) = func {
        let mut s = Vec::new();
        collect_semantic_from_expr(&f.body, &mut s);
        s
    } else {
        Vec::new()
    };

    let tokens = match lex_with_offset(body_slice, body_start) {
        Ok(t) => t,
        Err(_) => {
            out.extend(semantic);
            out.sort_by_key(|(s, _)| s.start);
            return out;
        }
    };

    out.extend(merge_semantic_lexical(semantic, tokens));
    out.sort_by_key(|(s, _)| s.start);
    out
}

/// Lexical-only highlighting for a single line (e.g. REPL input). No full parse.
/// Handles comment from `#` to end of line, then lexes the code part.
pub fn highlight_line(line: &str) -> Vec<(Span, HighlightKind)> {
    let mut out = Vec::new();
    let bytes = line.as_bytes();
    let comment_start = bytes.iter().position(|&b| b == b'#');
    let (code_end, comment_span) = match comment_start {
        Some(i) => {
            out.push((Span::new(i, line.len()), HighlightKind::Comment));
            (i, true)
        }
        None => (line.len(), false),
    };
    let code = &line[..code_end];
    if code.trim().is_empty() && comment_span {
        out.sort_by_key(|(s, _)| s.start);
        return out;
    }
    if let Ok(tokens) = lex_with_offset(code, 0) {
        for (tok, span) in tokens {
            out.push((span, token_kind_to_highlight(&tok)));
        }
    }
    out.sort_by_key(|(s, _)| s.start);
    out
}

/// ANSI color codes for terminal output (basic 16-color).
fn ansi_for_kind(kind: HighlightKind) -> &'static str {
    match kind {
        HighlightKind::Keyword => "\x1b[1;34m",   // bold blue
        HighlightKind::Variable => "\x1b[32m",    // green
        HighlightKind::Function => "\x1b[35m",    // magenta
        HighlightKind::Type => "\x1b[36m",        // cyan
        HighlightKind::Number => "\x1b[33m",      // yellow
        HighlightKind::Constant => "\x1b[1;33m",  // bold yellow
        HighlightKind::Operator => "\x1b[37m",    // white
        HighlightKind::Punctuation => "\x1b[90m", // bright black
        HighlightKind::Comment => "\x1b[2;90m",   // dim gray
    }
}
const ANSI_RESET: &str = "\x1b[0m";

/// Render a line with ANSI escape codes for syntax highlighting. Spans must be sorted by start.
pub fn spans_to_ansi(line: &str, spans: &[(Span, HighlightKind)]) -> String {
    if spans.is_empty() {
        return line.to_string();
    }
    let mut out = String::with_capacity(line.len() + spans.len() * 24);
    let mut i = 0;
    while i < line.len() {
        let c = line[i..].chars().next().unwrap_or('\0');
        let end = i + c.len_utf8();
        for (span, kind) in spans.iter() {
            if i >= span.start && i < span.end {
                if i == span.start {
                    out.push_str(ansi_for_kind(*kind));
                }
                break;
            }
        }
        out.push_str(&line[i..end]);
        for (span, _) in spans.iter() {
            if end == span.end {
                out.push_str(ANSI_RESET);
                break;
            }
        }
        i = end;
    }
    out
}

/// Convenience: highlight a single line and return it with ANSI colors for terminal.
pub fn highlight_line_ansi(line: &str) -> String {
    let spans = highlight_line(line);
    spans_to_ansi(line, &spans)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parse::parse_dsl;

    #[test]
    fn comment_span() {
        let src = "||\n# hello\n1";
        let func = parse_dsl(src).unwrap();
        let h = highlight(src, Some(&func));
        let comments: Vec<_> = h
            .iter()
            .filter(|(_, k)| *k == HighlightKind::Comment)
            .collect();
        assert_eq!(comments.len(), 1);
        assert_eq!(&src[comments[0].0.start..comments[0].0.end], "# hello");
    }

    #[test]
    fn header_param_and_type() {
        let src = "|a: i64, b: f32|\n1";
        let func = parse_dsl(src).unwrap();
        let h = highlight(src, Some(&func));
        let vars: Vec<_> = h
            .iter()
            .filter(|(_, k)| *k == HighlightKind::Variable)
            .collect();
        let types: Vec<_> = h
            .iter()
            .filter(|(_, k)| *k == HighlightKind::Type)
            .collect();
        assert!(vars.len() >= 2);
        assert!(types.len() >= 2);
    }

    #[test]
    fn variable_and_number_in_body() {
        let src = "|x: i64|\nx + 1";
        let func = parse_dsl(src).unwrap();
        let h = highlight(src, Some(&func));
        let has_var = h.iter().any(|(_, k)| *k == HighlightKind::Variable);
        let has_num = h.iter().any(|(_, k)| *k == HighlightKind::Number);
        assert!(has_var);
        assert!(has_num);
    }

    #[test]
    fn to_semantic_tokens_shape() {
        let src = "|n: i64|\nn + 1";
        let func = parse_dsl(src).unwrap();
        let spans = highlight(src, Some(&func));
        let tokens = to_semantic_tokens(&spans, src);
        assert_eq!(tokens.len() % 5, 0);
        assert!(!tokens.is_empty());
    }

    #[test]
    fn line_starts_and_byte_offset() {
        let src = "ab\ncd\ne";
        let starts = line_starts(src);
        assert_eq!(starts, [0, 3, 6]);
        assert_eq!(byte_offset_to_line_col(0, &starts), (0, 0));
        assert_eq!(byte_offset_to_line_col(2, &starts), (0, 2));
        assert_eq!(byte_offset_to_line_col(3, &starts), (1, 0));
        assert_eq!(byte_offset_to_line_col(5, &starts), (1, 2));
        assert_eq!(byte_offset_to_line_col(6, &starts), (2, 0));
    }

    #[test]
    fn highlight_line_lexical() {
        let line = "x + 1";
        let spans = highlight_line(line);
        let kinds: Vec<_> = spans.iter().map(|(_, k)| k).collect();
        assert!(kinds.contains(&&HighlightKind::Variable));
        assert!(kinds.contains(&&HighlightKind::Operator));
        assert!(kinds.contains(&&HighlightKind::Number));
    }

    #[test]
    fn highlight_line_ansi_contains_reset() {
        let line = "if true then 1 else 0";
        let ansi = highlight_line_ansi(line);
        assert!(ansi.contains("\x1b[0m"));
        assert!(ansi.contains("if"));
    }
}
