import Foundation
import AppKit
import Security

struct ConnectionProfile: Identifiable, Codable, Hashable {
    var id = UUID()
    var name = "Local MySQL"
    var host = "127.0.0.1"
    var port = 3306
    var username = "root"
    var database = ""
}
struct QueryResult {
    var columns: [String] = []
    var rows: [[String]] = []
    var message = ""
}

enum PasswordStore {
    static func read(_ id: UUID) -> String {
        let query: [String: Any] = [kSecClass as String: kSecClassGenericPassword,
                                    kSecAttrService as String: "WorkbenchReplacement.MySQL",
                                    kSecAttrAccount as String: id.uuidString,
                                    kSecReturnData as String: true,
                                    kSecMatchLimit as String: kSecMatchLimitOne]
        var value: CFTypeRef?
        guard SecItemCopyMatching(query as CFDictionary, &value) == errSecSuccess,
              let data = value as? Data else { return "" }
        return String(data: data, encoding: .utf8) ?? ""
    }
    static func save(_ password: String, id: UUID) {
        let query: [String: Any] = [kSecClass as String: kSecClassGenericPassword,
                                    kSecAttrService as String: "WorkbenchReplacement.MySQL",
                                    kSecAttrAccount as String: id.uuidString]
        SecItemDelete(query as CFDictionary)
        guard !password.isEmpty else { return }
        var attributes = query
        attributes[kSecValueData as String] = Data(password.utf8)
        SecItemAdd(attributes as CFDictionary, nil)
    }
}


struct WorkbenchQueryTab: Identifiable {
    let id: UUID
    var title: String
    var sql: String
    var result: QueryResult
    var fileURL: URL?

    init(
        id: UUID = UUID(),
        title: String,
        sql: String = "",
        result: QueryResult = QueryResult(),
        fileURL: URL? = nil
    ) {
        self.id = id
        self.title = title
        self.sql = sql
        self.result = result
        self.fileURL = fileURL
    }
}

struct SQLHistoryEntry: Identifiable, Codable {
    var id = UUID()
    var executedAt = Date()
    var connectionName: String
    var database: String
    var sql: String
    var succeeded: Bool
    var detail: String
}

struct SavedSQLQuery: Identifiable, Codable {
    var id = UUID()
    var name: String
    var sql: String
}

@MainActor
final class DatabaseModel: ObservableObject {
    @Published var profiles: [ConnectionProfile] = []
    @Published private(set) var queryHistory: [SQLHistoryEntry] = []
    @Published private(set) var savedQueries: [SavedSQLQuery] = []
    private let historyKey = "WorkbenchReplacement.queryHistory.v1"
    private let savedQueriesKey = "WorkbenchReplacement.savedQueries.v1"
    private let maxHistoryEntries = 200
    private let recentScriptsKey = "WorkbenchReplacement.recentSQLScripts.v1"
    @Published private(set) var recentSQLScripts: [String] = []

    private func recordRecentSQLScript(_ url: URL) {
        let path = url.standardizedFileURL.path
        recentSQLScripts.removeAll { $0 == path }
        recentSQLScripts.insert(path, at: 0)
        recentSQLScripts = Array(recentSQLScripts.prefix(12))
        UserDefaults.standard.set(recentSQLScripts, forKey: recentScriptsKey)
    }

    func clearRecentSQLScripts() {
        recentSQLScripts = []
        UserDefaults.standard.removeObject(forKey: recentScriptsKey)
    }

    func openRecentSQLScript(at path: String) {
        let url = URL(fileURLWithPath: path)
        do {
            let contents = try String(contentsOf: url, encoding: .utf8)
            saveActiveQuery()
            let tab = WorkbenchQueryTab(title: url.lastPathComponent, sql: contents, fileURL: url)
            queryTabs.append(tab)
            activateQuery(tab.id)
            recordRecentSQLScript(url)
            status = "Opened \(url.lastPathComponent)"
            NotificationCenter.default.post(name: .workbenchRevealSQLEditor, object: nil)
        } catch {
            status = "Could not open \(url.lastPathComponent): \(error.localizedDescription)"
        }
    }


    func saveCurrentQuery(named name: String) {
        let clean = name.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !clean.isEmpty, !sql.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
        savedQueries.insert(SavedSQLQuery(name: clean, sql: sql), at: 0)
        persistSavedQueries()
    }

    func openSavedQuery(_ entry: SavedSQLQuery) {
        newQueryTab()
        sql = entry.sql
        if let id = activeQueryID, let index = queryTabs.firstIndex(where: { $0.id == id }) {
            queryTabs[index].title = entry.name
            queryTabs[index].sql = entry.sql
        }
    }

    func deleteSavedQuery(_ id: UUID) {
        savedQueries.removeAll { $0.id == id }
        persistSavedQueries()
    }

    func clearQueryHistory() {
        queryHistory.removeAll()
        UserDefaults.standard.removeObject(forKey: historyKey)
    }

    private func persistSavedQueries() {
        if let data = try? JSONEncoder().encode(savedQueries) {
            UserDefaults.standard.set(data, forKey: savedQueriesKey)
        }
    }

    private func recordQuery(_ statement: String, profile: ConnectionProfile, database: String, succeeded: Bool, detail: String) {
        queryHistory.insert(SQLHistoryEntry(connectionName: profile.name, database: database, sql: statement, succeeded: succeeded, detail: detail), at: 0)
        if queryHistory.count > maxHistoryEntries {
            queryHistory.removeLast(queryHistory.count - maxHistoryEntries)
        }
        if let data = try? JSONEncoder().encode(queryHistory) {
            UserDefaults.standard.set(data, forKey: historyKey)
        }
    }

    @Published var selectedID: UUID?
    @Published var queryTabs: [WorkbenchQueryTab] = [
        WorkbenchQueryTab(
            title: "Query 1",
            sql: "SELECT VERSION() AS server_version;"
        )
    ]
    @Published var activeQueryID: UUID?
    @Published var sql = "SELECT VERSION() AS server_version;"
    @Published var result = QueryResult()
    @Published var databases: [String] = []
    @Published var schemaTables: [String: [String]] = [:]
    @Published var schemaViews: [String: [String]] = [:]
    @Published var schemaProcedures: [String: [String]] = [:]
    @Published var schemaFunctions: [String: [String]] = [:]
    @Published var schemaColumns: [String: [String]] = [:]
    // Introspection is a connection-scoped, RAM-only snapshot of table structures.
    struct IntrospectedColumn: Identifiable {
        var id: String { "\(table).\(name)" }
        let table: String
        let name: String
        let dataType: String
        let nullable: Bool
        let key: String
        let defaultValue: String?
        let extra: String
    }
    struct SchemaIntrospection {
        let connectionID: UUID
        let database: String
        let loadedAt: Date
        let tables: [String]
        let columnsByTable: [String: [IntrospectedColumn]]
    }
    @Published private(set) var introspections: [String: SchemaIntrospection] = [:]
    @Published private(set) var introspectionLoading: Set<String> = []
    private func introspectionKey(_ database: String, _ connectionID: UUID) -> String {
        "\(connectionID.uuidString):\(database)"
    }
    func introspection(for database: String) -> SchemaIntrospection? {
        guard let profile = selected else { return nil }
        return introspections[introspectionKey(database, profile.id)]
    }
    func introspect(database: String) {
        guard let profile = selected else { status = "Select a connection"; return }
        let key = introspectionKey(database, profile.id)
        guard !introspectionLoading.contains(key) else { return }
        introspectionLoading.insert(key)
        status = "Introspecting \(database)…"
        let password = PasswordStore.read(profile.id)
        let escaped = database.replacingOccurrences(of: "'", with: "''")
        Task {
            defer { introspectionLoading.remove(key) }
            do {
                let tableResult = try await NativeMySQLRunner.run(
                    profile: profile, database: database, password: password,
                    sql: "SELECT TABLE_NAME FROM information_schema.TABLES WHERE TABLE_SCHEMA = '\(escaped)' AND TABLE_TYPE = 'BASE TABLE' ORDER BY TABLE_NAME"
                )
                let columnResult = try await NativeMySQLRunner.run(
                    profile: profile, database: database, password: password,
                    sql: "SELECT TABLE_NAME, COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_KEY, COLUMN_DEFAULT, EXTRA FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '\(escaped)' ORDER BY TABLE_NAME, ORDINAL_POSITION"
                )
                let tables = tableResult.rows.compactMap { $0.first }
                var columns: [String: [IntrospectedColumn]] = [:]
                for row in columnResult.rows where row.count >= 7 {
                    let column = IntrospectedColumn(
                        table: row[0], name: row[1], dataType: row[2],
                        nullable: row[3] == "YES", key: row[4],
                        defaultValue: row[5] == "NULL" ? nil : row[5], extra: row[6]
                    )
                    columns[row[0], default: []].append(column)
                }
                // Ignore results from a connection that has been removed while loading.
                guard profiles.contains(where: { $0.id == profile.id }) else { return }
                introspections[key] = SchemaIntrospection(
                    connectionID: profile.id, database: database,
                    loadedAt: Date(), tables: tables, columnsByTable: columns
                )
                schemaTables[database] = tables
                schemaColumns[database] = Array(Set(columns.values.flatMap { $0 }.flatMap {
                    [$0.name, "\($0.table).\($0.name)"]
                })).sorted()
                status = "Introspection complete: \(tables.count) tables, \(columns.values.reduce(0) { $0 + $1.count }) columns in \(database) (RAM)"
            } catch {
                status = "Introspection failed for \(database): \(error.localizedDescription)"
            }
        }
    }
    @Published var selectedDatabase = ""
    @Published var queryDurationText = ""
    @Published var isRunning = false
    @Published var status = "Not connected"
    private let storageKey = "WorkbenchReplacement.connectionProfiles.native.v1"

    init() {
        recentSQLScripts = UserDefaults.standard.stringArray(forKey: recentScriptsKey) ?? []
        if let data = UserDefaults.standard.data(forKey: historyKey),
           let entries = try? JSONDecoder().decode([SQLHistoryEntry].self, from: data) {
            queryHistory = Array(entries.prefix(maxHistoryEntries))
        }
        if let data = UserDefaults.standard.data(forKey: savedQueriesKey),
           let entries = try? JSONDecoder().decode([SavedSQLQuery].self, from: data) {
            savedQueries = entries
        }
        if let data = UserDefaults.standard.data(forKey: storageKey),
           let decoded = try? JSONDecoder().decode([ConnectionProfile].self, from: data) { profiles = decoded }
        if profiles.isEmpty { profiles = [ConnectionProfile()] }
        selectedID = profiles.first?.id
        activeQueryID = queryTabs.first?.id
    }

    func newQueryTab() {
        saveActiveQuery()
        let tab = WorkbenchQueryTab(
            title: "Query \(queryTabs.count + 1)"
        )
        queryTabs.append(tab)
        activeQueryID = tab.id
        sql = tab.sql
        result = tab.result
    }

    func saveActiveQuery() {
        guard let id = activeQueryID,
              let index = queryTabs.firstIndex(where: { $0.id == id })
        else { return }

        queryTabs[index].sql = sql
        queryTabs[index].result = result
    }

    func activateQuery(_ id: UUID) {
        saveActiveQuery()

        guard let tab = queryTabs.first(where: { $0.id == id })
        else { return }

        activeQueryID = tab.id
        sql = tab.sql
        result = tab.result
    }

    func closeQuery(_ id: UUID) {
        saveActiveQuery()
        queryTabs.removeAll { $0.id == id }

        if queryTabs.isEmpty {
            let tab = WorkbenchQueryTab(title: "Query 1")
            queryTabs.append(tab)
        }

        if activeQueryID == id {
            activeQueryID = nil
            activateQuery(queryTabs[0].id)
        }
    }

    func openSQLScript() {
        let panel = NSOpenPanel()
        panel.allowedContentTypes = [.plainText]
        panel.allowsOtherFileTypes = true
        panel.canChooseDirectories = false

        guard panel.runModal() == .OK,
              let url = panel.url,
              let contents = try? String(contentsOf: url, encoding: .utf8)
        else { return }

        saveActiveQuery()

        let tab = WorkbenchQueryTab(
            title: url.lastPathComponent,
            sql: contents,
            fileURL: url
        )

        queryTabs.append(tab)
        activateQuery(tab.id)
        recordRecentSQLScript(url)
        NotificationCenter.default.post(name: .workbenchRevealSQLEditor, object: nil)
        status = "Opened \\(url.lastPathComponent)"
    }

    func saveSQLScript(asNewFile: Bool = false) {
        saveActiveQuery()

        guard let id = activeQueryID,
              let index = queryTabs.firstIndex(where: { $0.id == id })
        else { return }

        var destination = asNewFile ? nil : queryTabs[index].fileURL

        if destination == nil {
            let panel = NSSavePanel()
            panel.nameFieldStringValue = "query.sql"
            panel.canCreateDirectories = true

            guard panel.runModal() == .OK else { return }
            destination = panel.url
        }

        guard let url = destination else { return }

        do {
            try sql.write(to: url, atomically: true, encoding: .utf8)
            queryTabs[index].fileURL = url
            queryTabs[index].title = url.lastPathComponent
            recordRecentSQLScript(url)
            status = "Saved \\(url.lastPathComponent)"
        } catch {
            status = "Save failed: \\(error.localizedDescription)"
        }
    }

    var selected: ConnectionProfile? { profiles.first { $0.id == selectedID } }
    func save(_ profile: ConnectionProfile, password: String) {
        if let i = profiles.firstIndex(where: { $0.id == profile.id }) { profiles[i] = profile }
        else { profiles.append(profile) }
        PasswordStore.save(password, id: profile.id)
        selectedID = profile.id
        persist()
    }
    private func persist() {
        if let data = try? JSONEncoder().encode(profiles) { UserDefaults.standard.set(data, forKey: storageKey) }
    }

    func formatSQL() {
        sql = SQLFormatter.uppercaseKeywords(sql)
    }


    func loadTables(database: String) {
        guard let profile = selected else {
            status = "Select a MySQL connection"
            return
        }

        let password = PasswordStore.read(profile.id)
        let escaped = database.replacingOccurrences(
            of: "`",
            with: "``"
        )

        status = "Loading tables from \(database)…"

        Task {
            do {
                let parsed = try await NativeMySQLRunner.run(
                    profile: profile,
                    database: database,
                    password: password,
                    sql: "SHOW FULL TABLES FROM `\(escaped)`"
                )

                var tables: [String] = []
                var views: [String] = []

                for row in parsed.rows {
                    guard let name = row.first else {
                        continue
                    }

                    if row.count > 1 && row[1] == "VIEW" {
                        views.append(name)
                    } else {
                        tables.append(name)
                    }
                }

                schemaTables[database] = tables
                schemaViews[database] = views

                // Routine metadata is separate from SHOW FULL TABLES.
                // Failure (e.g. insufficient INFORMATION_SCHEMA privileges)
                // must not prevent browsing tables and views.
                do {
                    let quotedSchema = database.replacingOccurrences(of: "'", with: "''")
                    let routineResult = try await NativeMySQLRunner.run(
                        profile: profile, database: database, password: password,
                        sql: "SELECT ROUTINE_NAME, ROUTINE_TYPE FROM information_schema.ROUTINES WHERE ROUTINE_SCHEMA = '\(quotedSchema)' ORDER BY ROUTINE_NAME"
                    )
                    schemaProcedures[database] = routineResult.rows.compactMap { row in
                        row.count >= 2 && row[1].uppercased() == "PROCEDURE" ? row[0] : nil
                    }
                    schemaFunctions[database] = routineResult.rows.compactMap { row in
                        row.count >= 2 && row[1].uppercased() == "FUNCTION" ? row[0] : nil
                    }
                } catch {
                    schemaProcedures[database] = []
                    schemaFunctions[database] = []
                }

                // Cache column names for SQL completion. The lookup is read-only,
                // scoped to the selected schema and uses the current connection.
                // Failures (e.g. metadata permissions) do not prevent table browsing.
                do {
                    let quotedSchema = database.replacingOccurrences(of: "'", with: "''")
                    let columnResult = try await NativeMySQLRunner.run(
                        profile: profile,
                        database: database,
                        password: password,
                        sql: "SELECT TABLE_NAME, COLUMN_NAME FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = '\(quotedSchema)' ORDER BY TABLE_NAME, ORDINAL_POSITION"
                    )
                    var names = Set<String>()
                    for row in columnResult.rows where row.count >= 2 {
                        names.insert(row[1])
                        names.insert("\(row[0]).\(row[1])")
                    }
                    schemaColumns[database] = names.sorted()
                } catch {
                    schemaColumns[database] = []
                }

                status = "Loaded \(tables.count) tables and \(views.count) views from \(database)"
            } catch {
                status = "Unable to load tables: \(error.localizedDescription)"
                result = QueryResult(message: error.localizedDescription)
            }
        }
    }



    func selectTable(database: String, table: String) {
        let db = database.replacingOccurrences(
            of: "`",
            with: "``"
        )

        let tbl = table.replacingOccurrences(
            of: "`",
            with: "``"
        )

        selectedDatabase = database

        let configuredLimit = UserDefaults.standard.integer(
            forKey: "preferences.defaultRowLimit"
        )

        let rowLimit = configuredLimit > 0
            ? min(configuredLimit, 100000)
            : 500

        sql = "SELECT * FROM `\(db)`.`\(tbl)` LIMIT \(rowLimit);"

        runSQL()
    }

    func loadDatabases() { execute("SHOW DATABASES;", listingDatabases: true) }

    func runSQL() {
        if UserDefaults.standard.object(
            forKey: "preferences.uppercaseKeywords"
        ) as? Bool ?? true {
            formatSQL()
        }
        execute(sql, listingDatabases: false)
    }
    func runSQL(statement: String) {
        guard !statement.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
        execute(statement, listingDatabases: false)
    }
    private func execute(_ statement: String, listingDatabases: Bool) {
        guard let profile = selected else { status = "Select a connection"; return }
        let statements = listingDatabases ? [statement] : SQLStatementSplitter.split(statement)
        guard !statements.isEmpty else { return }
        isRunning = true
        if !listingDatabases { queryDurationText = "Running…" }
        let queryStarted = Date()
        status = "Executing…"
        let password = PasswordStore.read(profile.id)
        Task {
            var activeDB = listingDatabases ? "" : (selectedDatabase.isEmpty ? profile.database : selectedDatabase)
            do {
                var lastResult = QueryResult()
                for item in statements {
                    // Each query is sent separately: COM_QUERY accepts one SQL
                    // statement, not an entire semicolon-separated editor buffer.
                    let parsed = try await NativeMySQLRunner.run(
                        profile: profile, database: activeDB, password: password, sql: item
                    )
                    lastResult = parsed
                    if let usedSchema = SQLStatementSplitter.usedDatabase(item) {
                        activeDB = usedSchema
                        selectedDatabase = usedSchema
                        // Metadata is cached in memory for native SQL completion.
                        loadTables(database: usedSchema)
                    }
                }
                if listingDatabases {
                    databases = lastResult.rows.compactMap { $0.first }
                    status = "Connected to \(profile.name)"
                } else {
                    result = lastResult
                    if let id = activeQueryID,
                       let index = queryTabs.firstIndex(where: { $0.id == id }) {
                        queryTabs[index].sql = sql
                        queryTabs[index].result = lastResult
                    }
                    status = "Executed \(statements.count) statement(s) (\(lastResult.rows.count) rows)"
                    recordQuery(statement, profile: profile, database: activeDB, succeeded: true, detail: status)
                }
            } catch {
                status = "Error"
                result = QueryResult(message: error.localizedDescription)
                if !listingDatabases {
                    recordQuery(statement, profile: profile, database: activeDB, succeeded: false, detail: error.localizedDescription)
                }
            }
            if !listingDatabases {
                queryDurationText = String(format: "%.3f s", Date().timeIntervalSince(queryStarted))
            }
            isRunning = false
        }
    }

}


// Splits editor scripts without splitting semicolons inside quoted values or comments.
private enum SQLStatementSplitter {
    static func split(_ source: String) -> [String] {
        let chars = Array(source)
        var output: [String] = []
        var current = ""
        var state = 0 // 0 normal, 1 single quote, 2 double quote, 3 identifier, 4 line comment, 5 block comment
        var i = 0
        while i < chars.count {
            let c = chars[i]
            let next = i + 1 < chars.count ? chars[i + 1] : "\0"
            if state == 4 {
                if c == "\n" || c == "\r" { state = 0; current.append(" ") }
                i += 1; continue
            }
            if state == 5 {
                if c == "*" && next == "/" { state = 0; i += 2; current.append(" ") }
                else { i += 1 }
                continue
            }
            if state == 0 {
                if c == "#" || (c == "-" && next == "-" && (i + 2 == chars.count || chars[i + 2].isWhitespace)) {
                    state = 4; i += c == "#" ? 1 : 2; continue
                }
                if c == "/" && next == "*" { state = 5; i += 2; continue }
                if c == "'" { state = 1 }
                else if c == "\"" { state = 2 }
                else if c == "`" { state = 3 }
                else if c == ";" {
                    let trimmed = current.trimmingCharacters(in: .whitespacesAndNewlines)
                    if !trimmed.isEmpty { output.append(trimmed) }
                    current = ""; i += 1; continue
                }
            } else {
                if c == "\\" && state != 3 && i + 1 < chars.count {
                    current.append(c); current.append(next); i += 2; continue
                }
                if (state == 1 && c == "'") || (state == 2 && c == "\"") || (state == 3 && c == "`") {
                    if next == c { current.append(c); current.append(next); i += 2; continue }
                    state = 0
                }
            }
            current.append(c); i += 1
        }
        let trimmed = current.trimmingCharacters(in: .whitespacesAndNewlines)
        if !trimmed.isEmpty { output.append(trimmed) }
        return output
    }

    static func usedDatabase(_ statement: String) -> String? {
        let pattern = #"(?is)^\s*USE\s+(?:`((?:``|[^`])+)`|([A-Za-z0-9_$]+))\s*$"#
        guard let regex = try? NSRegularExpression(pattern: pattern),
              let match = regex.firstMatch(in: statement, range: NSRange(statement.startIndex..., in: statement)) else { return nil }
        for index in [1, 2] where match.range(at: index).location != NSNotFound {
            if let range = Range(match.range(at: index), in: statement) {
                return String(statement[range]).replacingOccurrences(of: "``", with: "`")
            }
        }
        return nil
    }
}

// MARK: - SQL Keyword Formatter

enum SQLFormatter {

    private static let keywords: Set<String> = [
        "SELECT", "FROM", "WHERE", "INSERT", "INTO", "VALUES",
        "UPDATE", "DELETE", "CREATE", "ALTER", "DROP", "TABLE",
        "DATABASE", "SCHEMA", "USE", "SHOW", "DESCRIBE", "EXPLAIN",
        "JOIN", "INNER", "LEFT", "RIGHT", "FULL", "OUTER",
        "ON", "AS", "AND", "OR", "NOT", "NULL", "IS", "IN",
        "BETWEEN", "LIKE", "ORDER", "BY", "GROUP", "HAVING",
        "LIMIT", "OFFSET", "ASC", "DESC", "DISTINCT", "ALL",
        "PRIMARY", "KEY", "FOREIGN", "REFERENCES", "INDEX",
        "UNIQUE", "DEFAULT", "AUTO_INCREMENT", "IF", "EXISTS",
        "CASCADE", "RESTRICT", "SET", "BEGIN", "COMMIT",
        "ROLLBACK", "TRANSACTION", "UNION", "CASE", "WHEN",
        "THEN", "ELSE", "END", "TRUE", "FALSE", "COUNT",
        "SUM", "AVG", "MIN", "MAX", "HAVING", "WITH",
        "RECURSIVE", "TRUNCATE", "REPLACE", "VIEW",
        "PROCEDURE", "FUNCTION", "TRIGGER", "GRANT",
        "REVOKE", "DISTINCT", "CURRENT_TIMESTAMP"
    ]

    static func uppercaseKeywords(_ sql: String) -> String {
        let chars = Array(sql)
        var output = ""
        var index = 0

        enum State {
            case normal
            case singleQuote
            case doubleQuote
            case backtick
            case lineComment
            case blockComment
        }

        var state = State.normal

        while index < chars.count {
            let ch = chars[index]
            let next: Character? = index + 1 < chars.count
                ? chars[index + 1] : nil

            switch state {
            case .normal:
                if ch == "'" {
                    state = .singleQuote
                    output.append(ch)
                    index += 1
                } else if ch == "\"" {
                    state = .doubleQuote
                    output.append(ch)
                    index += 1
                } else if ch == "`" {
                    state = .backtick
                    output.append(ch)
                    index += 1
                } else if ch == "-" && next == "-" &&
                          (index + 2 == chars.count ||
                           chars[index + 2].isWhitespace) {
                    state = .lineComment
                    output += "--"
                    index += 2
                } else if ch == "#" {
                    state = .lineComment
                    output.append(ch)
                    index += 1
                } else if ch == "/" && next == "*" {
                    state = .blockComment
                    output += "/*"
                    index += 2
                } else if ch.isLetter || ch == "_" {
                    let start = index
                    while index < chars.count &&
                          (chars[index].isLetter ||
                           chars[index].isNumber ||
                           chars[index] == "_" ||
                           chars[index] == "$") {
                        index += 1
                    }

                    let word = String(chars[start..<index])
                    let upper = word.uppercased()

                    output += keywords.contains(upper) ? upper : word
                } else {
                    output.append(ch)
                    index += 1
                }

            case .singleQuote:
                output.append(ch)
                index += 1
                if ch == "\\" && index < chars.count {
                    output.append(chars[index])
                    index += 1
                } else if ch == "'" {
                    if index < chars.count && chars[index] == "'" {
                        output.append(chars[index])
                        index += 1
                    } else {
                        state = .normal
                    }
                }

            case .doubleQuote:
                output.append(ch)
                index += 1
                if ch == "\\" && index < chars.count {
                    output.append(chars[index])
                    index += 1
                } else if ch == "\"" {
                    if index < chars.count && chars[index] == "\"" {
                        output.append(chars[index])
                        index += 1
                    } else {
                        state = .normal
                    }
                }

            case .backtick:
                output.append(ch)
                index += 1
                if ch == "`" {
                    if index < chars.count && chars[index] == "`" {
                        output.append(chars[index])
                        index += 1
                    } else {
                        state = .normal
                    }
                }

            case .lineComment:
                output.append(ch)
                index += 1
                if ch == "\n" {
                    state = .normal
                }

            case .blockComment:
                output.append(ch)
                index += 1
                if ch == "*" && index < chars.count &&
                   chars[index] == "/" {
                    output.append("/")
                    index += 1
                    state = .normal
                }
            }
        }

        return output
    }
}