//
//  SyntaxHighlighter.swift
//  oAI
//
//  Keyword-based syntax highlighting using AttributedString
//

import SwiftUI

struct SyntaxHighlighter {

    // MARK: - Token Colors (dark theme)

    static let keywordColor = Color(hex: "#569cd6")      // Blue
    static let stringColor = Color(hex: "#ce9178")        // Orange
    static let commentColor = Color(hex: "#6a9955")       // Green
    static let numberColor = Color(hex: "#b5cea8")        // Light green
    static let typeColor = Color(hex: "#4ec9b0")          // Teal
    static let functionColor = Color(hex: "#dcdcaa")      // Yellow
    static let defaultColor = Color(hex: "#d4d4d4")       // Light gray
    static let punctuationColor = Color(hex: "#808080")   // Gray

    // MARK: - Language Keywords

    private static let keywords: [String: Set<String>] = [
        "swift": ["import", "func", "class", "struct", "enum", "protocol", "extension",
                  "var", "let", "if", "else", "guard", "switch", "case", "default",
                  "for", "while", "repeat", "return", "break", "continue", "throw",
                  "throws", "try", "catch", "do", "async", "await", "in", "where",
                  "self", "Self", "super", "init", "deinit", "nil", "true", "false",
                  "public", "private", "internal", "fileprivate", "open", "static",
                  "override", "mutating", "weak", "unowned", "lazy", "some", "any",
                  "typealias", "associatedtype", "inout", "as", "is", "defer"],
        "python": ["import", "from", "def", "class", "if", "elif", "else", "for",
                   "while", "return", "yield", "break", "continue", "pass", "raise",
                   "try", "except", "finally", "with", "as", "lambda", "and", "or",
                   "not", "in", "is", "True", "False", "None", "self", "async", "await",
                   "global", "nonlocal", "del", "assert", "print"],
        "javascript": ["function", "const", "let", "var", "if", "else", "for", "while",
                       "do", "switch", "case", "default", "return", "break", "continue",
                       "throw", "try", "catch", "finally", "class", "extends", "new",
                       "this", "super", "import", "export", "from", "async", "await",
                       "yield", "of", "in", "typeof", "instanceof", "delete", "void",
                       "true", "false", "null", "undefined", "debugger"],
        "typescript": ["function", "const", "let", "var", "if", "else", "for", "while",
                       "do", "switch", "case", "default", "return", "break", "continue",
                       "throw", "try", "catch", "finally", "class", "extends", "new",
                       "this", "super", "import", "export", "from", "async", "await",
                       "yield", "of", "in", "typeof", "instanceof", "delete", "void",
                       "true", "false", "null", "undefined", "type", "interface",
                       "enum", "implements", "abstract", "readonly", "as", "keyof",
                       "namespace", "declare", "module"],
        "go": ["package", "import", "func", "type", "struct", "interface", "var",
               "const", "if", "else", "for", "range", "switch", "case", "default",
               "return", "break", "continue", "go", "defer", "select", "chan",
               "map", "make", "new", "append", "len", "cap", "nil", "true", "false",
               "fallthrough", "goto"],
        "rust": ["fn", "let", "mut", "const", "if", "else", "match", "for", "while",
                "loop", "return", "break", "continue", "struct", "enum", "impl",
                "trait", "pub", "use", "mod", "crate", "self", "super", "as", "in",
                "ref", "move", "async", "await", "where", "type", "dyn", "unsafe",
                "extern", "true", "false", "Some", "None", "Ok", "Err"],
        "java": ["class", "interface", "enum", "extends", "implements", "import",
                 "package", "public", "private", "protected", "static", "final",
                 "abstract", "void", "int", "long", "double", "float", "boolean",
                 "char", "byte", "short", "if", "else", "for", "while", "do",
                 "switch", "case", "default", "return", "break", "continue",
                 "throw", "throws", "try", "catch", "finally", "new", "this",
                 "super", "null", "true", "false", "synchronized", "volatile"],
        "c": ["if", "else", "for", "while", "do", "switch", "case", "default",
              "return", "break", "continue", "goto", "typedef", "struct", "union",
              "enum", "const", "static", "extern", "volatile", "register", "auto",
              "void", "int", "long", "short", "char", "float", "double", "unsigned",
              "signed", "sizeof", "NULL", "include", "define", "ifdef", "ifndef",
              "endif", "pragma"],
        "cpp": ["if", "else", "for", "while", "do", "switch", "case", "default",
                "return", "break", "continue", "goto", "typedef", "struct", "union",
                "enum", "const", "static", "extern", "volatile", "class", "public",
                "private", "protected", "virtual", "override", "template", "typename",
                "namespace", "using", "new", "delete", "throw", "try", "catch",
                "nullptr", "true", "false", "auto", "constexpr", "inline",
                "void", "int", "long", "short", "char", "float", "double", "bool",
                "include", "define", "ifdef", "ifndef", "endif"],
        "sql": ["SELECT", "FROM", "WHERE", "INSERT", "INTO", "VALUES", "UPDATE",
                "SET", "DELETE", "CREATE", "TABLE", "ALTER", "DROP", "INDEX",
                "JOIN", "LEFT", "RIGHT", "INNER", "OUTER", "ON", "AND", "OR",
                "NOT", "NULL", "IS", "IN", "LIKE", "BETWEEN", "EXISTS", "AS",
                "ORDER", "BY", "GROUP", "HAVING", "LIMIT", "OFFSET", "UNION",
                "DISTINCT", "COUNT", "SUM", "AVG", "MAX", "MIN", "PRIMARY",
                "KEY", "FOREIGN", "REFERENCES", "CASCADE", "CONSTRAINT",
                "select", "from", "where", "insert", "into", "values", "update",
                "set", "delete", "create", "table", "alter", "drop", "index",
                "join", "left", "right", "inner", "outer", "on", "and", "or",
                "not", "null", "is", "in", "like", "between", "exists", "as",
                "order", "by", "group", "having", "limit", "offset", "union",
                "distinct", "primary", "key", "foreign", "references"],
        "shell": ["if", "then", "else", "elif", "fi", "for", "while", "do", "done",
                  "case", "esac", "function", "return", "exit", "echo", "export",
                  "local", "readonly", "shift", "set", "unset", "source", "eval",
                  "exec", "cd", "pwd", "ls", "cp", "mv", "rm", "mkdir", "cat",
                  "grep", "sed", "awk", "find", "xargs", "pipe", "true", "false",
                  "in", "sudo", "chmod", "chown"],
        "html": ["html", "head", "body", "div", "span", "p", "a", "img", "ul", "ol",
                "li", "table", "tr", "td", "th", "form", "input", "button", "select",
                "option", "textarea", "script", "style", "link", "meta", "title",
                "header", "footer", "nav", "main", "section", "article", "aside",
                "class", "id", "href", "src", "alt", "type", "value", "name"],
        "css": ["color", "background", "margin", "padding", "border", "font",
               "display", "position", "width", "height", "top", "left", "right",
               "bottom", "flex", "grid", "align", "justify", "transform", "transition",
               "animation", "opacity", "overflow", "z-index", "important",
               "none", "block", "inline", "absolute", "relative", "fixed", "sticky"],
        "ruby": ["def", "end", "class", "module", "if", "elsif", "else", "unless",
                "while", "until", "for", "do", "begin", "rescue", "ensure", "raise",
                "return", "yield", "block_given?", "require", "include", "extend",
                "attr_accessor", "attr_reader", "attr_writer", "self", "super",
                "nil", "true", "false", "puts", "print", "lambda", "proc"],
    ]

    // MARK: - Comment Styles

    private static let lineCommentPrefixes: [String: String] = [
        "swift": "//", "python": "#", "javascript": "//", "typescript": "//",
        "go": "//", "rust": "//", "java": "//", "c": "//", "cpp": "//",
        "shell": "#", "bash": "#", "ruby": "#", "yaml": "#", "toml": "#",
    ]

    // MARK: - Language Aliases

    private static let languageAliases: [String: String] = [
        "js": "javascript", "ts": "typescript", "py": "python",
        "sh": "shell", "bash": "shell", "zsh": "shell",
        "c++": "cpp", "objective-c": "c", "objc": "c",
        "yml": "yaml", "md": "markdown", "rb": "ruby",
        "h": "c", "hpp": "cpp", "m": "c",
    ]

    // MARK: - Public API

    static func highlight(code: String, language: String?) -> AttributedString {
        let lang = resolveLanguage(language)
        let langKeywords = keywords[lang] ?? Set()
        let commentPrefix = lineCommentPrefixes[lang]

        var result = AttributedString()
        let lines = code.components(separatedBy: "\n")

        for (lineIndex, line) in lines.enumerated() {
            let highlightedLine = highlightLine(line, keywords: langKeywords, commentPrefix: commentPrefix, language: lang)
            result.append(highlightedLine)
            if lineIndex < lines.count - 1 {
                result.append(AttributedString("\n"))
            }
        }

        return result
    }

    // MARK: - Private

    private static func resolveLanguage(_ lang: String?) -> String {
        guard let lang = lang?.lowercased().trimmingCharacters(in: .whitespaces), !lang.isEmpty else {
            return ""
        }
        return languageAliases[lang] ?? lang
    }

    private static func highlightLine(_ line: String, keywords: Set<String>, commentPrefix: String?, language: String) -> AttributedString {
        // Check for line comments
        if let prefix = commentPrefix {
            let trimmed = line.trimmingCharacters(in: .whitespaces)
            if trimmed.hasPrefix(prefix) {
                var attr = AttributedString(line)
                attr.foregroundColor = commentColor
                return attr
            }
        }

        // Tokenize and colorize
        var result = AttributedString()
        var i = line.startIndex

        while i < line.endIndex {
            let c = line[i]

            // String literals
            if c == "\"" || c == "'" || c == "`" {
                let (strAttr, newIndex) = consumeString(line, from: i, quote: c)
                result.append(strAttr)
                i = newIndex
                continue
            }

            // Block comment start
            if c == "/" && line.index(after: i) < line.endIndex && line[line.index(after: i)] == "*" {
                // Consume rest of line as comment (simplified — no multi-line tracking)
                let rest = String(line[i...])
                var attr = AttributedString(rest)
                attr.foregroundColor = commentColor
                result.append(attr)
                return result
            }

            // Numbers
            if c.isNumber && (i == line.startIndex || !line[line.index(before: i)].isLetter) {
                let (numAttr, newIndex) = consumeNumber(line, from: i)
                result.append(numAttr)
                i = newIndex
                continue
            }

            // Words (identifiers/keywords)
            if c.isLetter || c == "_" || c == "@" || c == "#" {
                let (wordAttr, newIndex) = consumeWord(line, from: i, keywords: keywords)
                result.append(wordAttr)
                i = newIndex
                continue
            }

            // Punctuation/operators
            var charAttr = AttributedString(String(c))
            charAttr.foregroundColor = defaultColor
            result.append(charAttr)
            i = line.index(after: i)
        }

        return result
    }

    private static func consumeString(_ line: String, from start: String.Index, quote: Character) -> (AttributedString, String.Index) {
        var i = line.index(after: start)
        var str = String(quote)

        while i < line.endIndex {
            let c = line[i]
            str.append(c)
            if c == "\\" && line.index(after: i) < line.endIndex {
                // Escaped character
                i = line.index(after: i)
                str.append(line[i])
                i = line.index(after: i)
                continue
            }
            i = line.index(after: i)
            if c == quote {
                break
            }
        }

        var attr = AttributedString(str)
        attr.foregroundColor = stringColor
        return (attr, i)
    }

    private static func consumeNumber(_ line: String, from start: String.Index) -> (AttributedString, String.Index) {
        var i = start
        var num = ""

        while i < line.endIndex && (line[i].isHexDigit || line[i] == "." || line[i] == "x" || line[i] == "X" || line[i] == "_") {
            num.append(line[i])
            i = line.index(after: i)
        }

        var attr = AttributedString(num)
        attr.foregroundColor = numberColor
        return (attr, i)
    }

    private static func consumeWord(_ line: String, from start: String.Index, keywords: Set<String>) -> (AttributedString, String.Index) {
        var i = start
        var word = ""

        while i < line.endIndex && (line[i].isLetter || line[i].isNumber || line[i] == "_" || line[i] == "@" || line[i] == "#") {
            word.append(line[i])
            i = line.index(after: i)
        }

        var attr = AttributedString(word)

        if keywords.contains(word) {
            attr.foregroundColor = keywordColor
        } else if word.first?.isUppercase == true && word.count > 1 {
            // Type-like identifier (capitalized)
            attr.foregroundColor = typeColor
        } else if i < line.endIndex && line[i] == "(" {
            // Function call
            attr.foregroundColor = functionColor
        } else {
            attr.foregroundColor = defaultColor
        }

        return (attr, i)
    }
}