import AppKit

/// A token category the highlighter recognizes. Kept deliberately small — a
/// handful of broad classes that read well across many languages beats a precise
/// per-language grammar we'd have to maintain.
enum SyntaxToken {
    case keyword, string, comment, number, type
}

/// The per-language rules the scanner runs on. One generic scanner (`SyntaxScanner`)
/// drives every language; a language is just this bag of rules. That keeps adding a
/// language to filling in keywords and comment markers rather than writing a parser.
struct LanguageSyntax {
    var keywords: Set<String> = []
    var lineComments: [String] = []
    var blockComment: (open: String, close: String)? = nil
    var stringDelimiters: [Character] = ["\""]
    /// Languages with `"""` heredoc-style multi-line strings (Swift, Python).
    var tripleQuoted = false
    /// Color capitalized identifiers as types. Pleasant in Swift/Rust/C-family,
    /// noisy in shells where ALL_CAPS env vars abound — off by default.
    var highlightTypes = false
    var highlightNumbers = true

    /// Whether this language has anything worth coloring; a `.plain` buffer skips
    /// the scan entirely.
    var isEmpty: Bool {
        keywords.isEmpty && lineComments.isEmpty && blockComment == nil
            && stringDelimiters.isEmpty && !tripleQuoted
    }
}

// MARK: - Language table

extension LanguageSyntax {
    /// No highlighting — plain monospaced text. Used for unknown extensions.
    static let plain = LanguageSyntax(stringDelimiters: [])

    static let swift = LanguageSyntax(
        keywords: [
            "associatedtype", "actor", "async", "await", "as", "break", "case", "catch",
            "class", "continue", "default", "defer", "deinit", "do", "else", "enum",
            "extension", "fallthrough", "false", "fileprivate", "final", "for", "func",
            "guard", "if", "import", "in", "init", "inout", "internal", "is", "lazy",
            "let", "mutating", "nil", "nonmutating", "open", "operator", "private",
            "protocol", "public", "repeat", "rethrows", "return", "self", "Self",
            "static", "struct", "subscript", "super", "switch", "throw", "throws",
            "true", "try", "typealias", "unowned", "var", "weak", "where", "while",
            "willSet", "didSet", "get", "set", "some", "any", "convenience", "required",
            "override", "indirect", "package",
        ],
        lineComments: ["//"], blockComment: ("/*", "*/"),
        tripleQuoted: true, highlightTypes: true)

    static let cLike = LanguageSyntax(
        keywords: [
            "auto", "bool", "break", "case", "catch", "char", "class", "const",
            "constexpr", "continue", "default", "delete", "do", "double", "else",
            "enum", "explicit", "extern", "false", "final", "float", "for", "friend",
            "goto", "if", "import", "inline", "int", "interface", "long", "namespace",
            "new", "nil", "null", "nullptr", "operator", "override", "package",
            "private", "protected", "public", "register", "return", "short", "signed",
            "sizeof", "static", "struct", "switch", "template", "this", "throw", "true",
            "try", "typedef", "typename", "union", "unsigned", "using", "virtual",
            "void", "volatile", "while", "var", "val", "fun", "let", "func", "implements",
            "extends", "abstract", "synchronized", "instanceof", "in", "is", "when",
        ],
        lineComments: ["//"], blockComment: ("/*", "*/"),
        stringDelimiters: ["\"", "'"], highlightTypes: true)

    static let javascript = LanguageSyntax(
        keywords: [
            "abstract", "any", "as", "async", "await", "boolean", "break", "case",
            "catch", "class", "const", "continue", "debugger", "declare", "default",
            "delete", "do", "else", "enum", "export", "extends", "false", "finally",
            "for", "from", "function", "get", "if", "implements", "import", "in",
            "instanceof", "interface", "let", "namespace", "new", "null", "number",
            "of", "private", "protected", "public", "readonly", "return", "set",
            "static", "string", "super", "switch", "this", "throw", "true", "try",
            "type", "typeof", "undefined", "var", "void", "while", "yield", "keyof",
        ],
        lineComments: ["//"], blockComment: ("/*", "*/"),
        stringDelimiters: ["\"", "'", "`"], highlightTypes: true)

    static let python = LanguageSyntax(
        keywords: [
            "and", "as", "assert", "async", "await", "break", "class", "continue",
            "def", "del", "elif", "else", "except", "False", "finally", "for", "from",
            "global", "if", "import", "in", "is", "lambda", "None", "nonlocal", "not",
            "or", "pass", "raise", "return", "True", "try", "while", "with", "yield",
            "match", "case", "self",
        ],
        lineComments: ["#"], stringDelimiters: ["\"", "'"],
        tripleQuoted: true, highlightTypes: true)

    static let ruby = LanguageSyntax(
        keywords: [
            "alias", "and", "begin", "break", "case", "class", "def", "defined?", "do",
            "else", "elsif", "end", "ensure", "false", "for", "if", "in", "module",
            "next", "nil", "not", "or", "redo", "rescue", "retry", "return", "self",
            "super", "then", "true", "undef", "unless", "until", "when", "while",
            "yield", "require", "attr_accessor", "attr_reader", "attr_writer",
        ],
        lineComments: ["#"], stringDelimiters: ["\"", "'"])

    static let go = LanguageSyntax(
        keywords: [
            "break", "case", "chan", "const", "continue", "default", "defer", "else",
            "fallthrough", "false", "for", "func", "go", "goto", "if", "import",
            "interface", "map", "nil", "package", "range", "return", "select", "struct",
            "switch", "true", "type", "var", "bool", "byte", "error", "int", "int64",
            "string", "uint", "rune", "float64",
        ],
        lineComments: ["//"], blockComment: ("/*", "*/"),
        stringDelimiters: ["\"", "`"], highlightTypes: true)

    static let rust = LanguageSyntax(
        keywords: [
            "as", "async", "await", "break", "const", "continue", "crate", "dyn",
            "else", "enum", "extern", "false", "fn", "for", "if", "impl", "in", "let",
            "loop", "match", "mod", "move", "mut", "pub", "ref", "return", "self",
            "Self", "static", "struct", "super", "trait", "true", "type", "unsafe",
            "use", "where", "while", "Some", "None", "Ok", "Err", "Option", "Result",
        ],
        lineComments: ["//"], blockComment: ("/*", "*/"),
        highlightTypes: true)

    static let shell = LanguageSyntax(
        keywords: [
            "if", "then", "else", "elif", "fi", "for", "while", "until", "do", "done",
            "case", "esac", "in", "function", "return", "select", "time", "export",
            "local", "readonly", "declare", "unset", "alias", "source", "exit", "set",
            "true", "false", "echo", "cd", "trap",
        ],
        lineComments: ["#"], stringDelimiters: ["\"", "'"])

    static let yaml = LanguageSyntax(
        keywords: ["true", "false", "null", "yes", "no", "on", "off", "~"],
        lineComments: ["#"], stringDelimiters: ["\"", "'"])

    static let json = LanguageSyntax(
        keywords: ["true", "false", "null"], stringDelimiters: ["\""])

    static let toml = LanguageSyntax(
        keywords: ["true", "false"],
        lineComments: ["#"], stringDelimiters: ["\"", "'"])

    /// Pick a language from a file's extension. Unknown extensions get `.plain`.
    static func forFile(_ url: URL?) -> LanguageSyntax {
        switch url?.pathExtension.lowercased() {
        case "swift": return .swift
        case "js", "jsx", "mjs", "cjs", "ts", "tsx": return .javascript
        case "py", "pyw", "pyi": return .python
        case "rb", "rake", "gemspec": return .ruby
        case "go": return .go
        case "rs": return .rust
        case "c", "h", "cpp", "cc", "cxx", "hpp", "hh", "m", "mm",
             "java", "kt", "kts", "cs", "swiftinterface":
            return .cLike
        case "sh", "bash", "zsh", "fish", "command": return .shell
        case "yml", "yaml": return .yaml
        case "json", "jsonc": return .json
        case "toml": return .toml
        default: return .plain
        }
    }
}

// MARK: - Scanner

/// A single-pass lexer shared by every language. At each position it tries, in
/// precedence order, the constructs that can start there (comment, string, number,
/// identifier) and emits a `(range, token)` for the ones worth coloring. Working in
/// `String.Index` space keeps the ranges correct for `NSRange(_:in:)` even with
/// multi-byte characters.
enum SyntaxScanner {
    static func tokens(in text: String, language: LanguageSyntax) -> [(Range<String.Index>, SyntaxToken)] {
        guard !language.isEmpty else { return [] }
        var out: [(Range<String.Index>, SyntaxToken)] = []
        let end = text.endIndex
        var i = text.startIndex

        // Does `text` start with `needle` at `idx`? Cheap manual prefix check so we
        // avoid allocating substrings in the hot loop.
        func matches(_ needle: String, at idx: String.Index) -> Bool {
            var a = idx
            for ch in needle {
                if a == end || text[a] != ch { return false }
                a = text.index(after: a)
            }
            return true
        }
        func advance(_ idx: inout String.Index, _ n: Int) {
            var k = 0
            while k < n && idx != end { idx = text.index(after: idx); k += 1 }
        }

        while i < end {
            let c = text[i]

            // Line comment — to end of line.
            if let lc = language.lineComments.first(where: { matches($0, at: i) }) {
                _ = lc
                let start = i
                while i < end && text[i] != "\n" { i = text.index(after: i) }
                out.append((start..<i, .comment))
                continue
            }

            // Block comment — to the closing delimiter (or end of file).
            if let bc = language.blockComment, matches(bc.open, at: i) {
                let start = i
                advance(&i, bc.open.count)
                while i < end && !matches(bc.close, at: i) { i = text.index(after: i) }
                advance(&i, bc.close.count)
                out.append((start..<i, .comment))
                continue
            }

            // Triple-quoted string — spans lines.
            if language.tripleQuoted && matches("\"\"\"", at: i) {
                let start = i
                advance(&i, 3)
                while i < end && !matches("\"\"\"", at: i) { i = text.index(after: i) }
                advance(&i, 3)
                out.append((start..<i, .string))
                continue
            }

            // Single-line string — honors backslash escapes, stops at end of line if
            // unterminated so a stray quote doesn't paint the rest of the file.
            if language.stringDelimiters.contains(c) {
                let start = i
                i = text.index(after: i)
                while i < end {
                    let ch = text[i]
                    if ch == "\\" {
                        i = text.index(after: i)
                        if i < end { i = text.index(after: i) }
                        continue
                    }
                    if ch == "\n" { break }
                    i = text.index(after: i)
                    if ch == c { break }
                }
                out.append((start..<i, .string))
                continue
            }

            // Number — a digit run plus the usual hex/float/exponent companions.
            if language.highlightNumbers && c.isNumber {
                let start = i
                while i < end, isNumberChar(text[i]) { i = text.index(after: i) }
                out.append((start..<i, .number))
                continue
            }

            // Identifier — a keyword, a capitalized type, or nothing to color.
            if c.isLetter || c == "_" {
                let start = i
                while i < end, let ch = Optional(text[i]), ch.isLetter || ch.isNumber || ch == "_" {
                    i = text.index(after: i)
                }
                let word = text[start..<i]
                if language.keywords.contains(String(word)) {
                    out.append((start..<i, .keyword))
                } else if language.highlightTypes, let f = word.first, f.isUppercase {
                    out.append((start..<i, .type))
                }
                continue
            }

            i = text.index(after: i)
        }
        return out
    }

    private static func isNumberChar(_ c: Character) -> Bool {
        c.isHexDigit || c == "." || c == "_" || c == "x" || c == "X"
            || c == "o" || c == "b" || c == "e" || c == "E" || c == "+" || c == "-"
    }
}

// MARK: - Colors

/// A resolved set of token colors. Each `ColorVisionMode` gets its own set so the
/// editor's palette stays distinguishable under the common color-vision deficiencies
/// — the same approach `AppPalette` takes for the app's categorical colors. Colors
/// still resolve per light/dark appearance.
struct SyntaxPalette {
    let keyword: NSColor
    let string: NSColor
    let comment: NSColor
    let number: NSColor
    let type: NSColor
    /// Default (un-tokenized) text color — matches `AppTheme.primaryText`.
    let plain: NSColor

    func color(_ token: SyntaxToken) -> NSColor {
        switch token {
        case .keyword: return keyword
        case .string: return string
        case .comment: return comment
        case .number: return number
        case .type: return type
        }
    }
}

/// Token colors per color-vision mode. The standard hues echo a calm Xcode-style
/// theme; the deficiency variants pull from the same color families `AppPalette`
/// uses (Okabe–Ito for red–green modes, a red/green/violet spread for tritanopia)
/// so the editor reads consistently with the rest of the UI. Comments stay a calm
/// gray and the plain text color is shared across modes.
enum SyntaxTheme {
    static func palette(for mode: ColorVisionMode) -> SyntaxPalette {
        switch mode {
        case .standard: return standard
        case .deuteranopia, .protanopia: return redGreenSafe
        case .tritanopia: return blueYellowSafe
        }
    }

    /// Shared across every mode — neutral, already-safe roles.
    private static let plain = dynamic(
        light: NSColor(srgbRed: 0.13, green: 0.13, blue: 0.15, alpha: 1),
        dark: NSColor(srgbRed: 0.80, green: 0.81, blue: 0.83, alpha: 1))
    private static let comment = dynamic(
        light: NSColor(srgbRed: 0.42, green: 0.48, blue: 0.52, alpha: 1),
        dark: NSColor(srgbRed: 0.45, green: 0.50, blue: 0.55, alpha: 1))

    /// Vivid Xcode-like default.
    private static let standard = SyntaxPalette(
        keyword: dynamic(light: NSColor(srgbRed: 0.68, green: 0.13, blue: 0.49, alpha: 1),
                         dark: NSColor(srgbRed: 1.00, green: 0.48, blue: 0.71, alpha: 1)),
        string: dynamic(light: NSColor(srgbRed: 0.77, green: 0.16, blue: 0.11, alpha: 1),
                        dark: NSColor(srgbRed: 0.99, green: 0.51, blue: 0.44, alpha: 1)),
        comment: comment,
        number: dynamic(light: NSColor(srgbRed: 0.11, green: 0.16, blue: 0.82, alpha: 1),
                        dark: NSColor(srgbRed: 0.85, green: 0.79, blue: 0.49, alpha: 1)),
        type: dynamic(light: NSColor(srgbRed: 0.06, green: 0.41, blue: 0.63, alpha: 1),
                      dark: NSColor(srgbRed: 0.40, green: 0.72, blue: 0.64, alpha: 1)),
        plain: plain)

    /// Deuteranopia/protanopia (red↔green weak). Okabe–Ito-derived: separation rides
    /// the blue–yellow axis and lightness — reddish-purple, vermillion, blue, green.
    private static let redGreenSafe = SyntaxPalette(
        keyword: dynamic(light: NSColor(srgbRed: 0.61, green: 0.24, blue: 0.51, alpha: 1),
                         dark: NSColor(srgbRed: 0.85, green: 0.55, blue: 0.72, alpha: 1)),
        string: dynamic(light: NSColor(srgbRed: 0.74, green: 0.32, blue: 0.00, alpha: 1),
                        dark: NSColor(srgbRed: 0.93, green: 0.55, blue: 0.24, alpha: 1)),
        comment: comment,
        number: dynamic(light: NSColor(srgbRed: 0.00, green: 0.40, blue: 0.66, alpha: 1),
                        dark: NSColor(srgbRed: 0.36, green: 0.71, blue: 0.92, alpha: 1)),
        type: dynamic(light: NSColor(srgbRed: 0.00, green: 0.52, blue: 0.42, alpha: 1),
                      dark: NSColor(srgbRed: 0.22, green: 0.76, blue: 0.60, alpha: 1)),
        plain: plain)

    /// Tritanopia (blue↔yellow weak; red↔green intact). Lean on red/green and
    /// lightness — magenta, red, green-teal, violet.
    private static let blueYellowSafe = SyntaxPalette(
        keyword: dynamic(light: NSColor(srgbRed: 0.80, green: 0.12, blue: 0.42, alpha: 1),
                         dark: NSColor(srgbRed: 0.95, green: 0.40, blue: 0.62, alpha: 1)),
        string: dynamic(light: NSColor(srgbRed: 0.74, green: 0.10, blue: 0.12, alpha: 1),
                        dark: NSColor(srgbRed: 0.92, green: 0.36, blue: 0.34, alpha: 1)),
        comment: comment,
        number: dynamic(light: NSColor(srgbRed: 0.00, green: 0.50, blue: 0.36, alpha: 1),
                        dark: NSColor(srgbRed: 0.22, green: 0.76, blue: 0.55, alpha: 1)),
        type: dynamic(light: NSColor(srgbRed: 0.34, green: 0.28, blue: 0.72, alpha: 1),
                      dark: NSColor(srgbRed: 0.62, green: 0.55, blue: 0.95, alpha: 1)),
        plain: plain)

    private static func dynamic(light: NSColor, dark: NSColor) -> NSColor {
        NSColor(name: nil) { appearance in
            appearance.bestMatch(from: [.aqua, .darkAqua]) == .darkAqua ? dark : light
        }
    }
}
