SQLWorkbench/WorkbenchReplacement/DatabaseModel.swift

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
    }
}

@MainActor
final class DatabaseModel: ObservableObject {
    @Published var profiles: [ConnectionProfile] = []
    @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 selectedDatabase = ""
    @Published var isRunning = false
    @Published var status = "Not connected"
    private let storageKey = "WorkbenchReplacement.connectionProfiles.native.v1"

    init() {
        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)
        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
            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

                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)
    }
    private func execute(_ statement: String, listingDatabases: Bool) {
        guard let profile = selected else { status = "Select a connection"; return }
        guard !statement.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
        isRunning = true
        status = "Connecting…"
        let activeDB = listingDatabases ? "" : (selectedDatabase.isEmpty ? profile.database : selectedDatabase)
        let password = PasswordStore.read(profile.id)
        Task {
            do {
                let parsed = try await NativeMySQLRunner.run(profile: profile, database: activeDB, password: password, sql: statement)
                if listingDatabases {
                    databases = parsed.rows.compactMap { $0.first }
                    status = "Connected to \(profile.name)"
                } else {
                    result = parsed
                    if let id = activeQueryID,
                       let index = queryTabs.firstIndex(where: { $0.id == id }) {
                        queryTabs[index].sql = sql
                        queryTabs[index].result = parsed
                    }
                    status = "Query completed (\(parsed.rows.count) rows)"
                }
            } catch {
                status = "Error"
                result = QueryResult(message: error.localizedDescription)
            }
            isRunning = false
        }
    }
}


// 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
    }
}