import Foundation
import CommonCrypto
import CoreFoundation
import Security

// Native MySQL wire protocol client. No external mysql executable is used.
// Initial scope: TCP, mysql_native_password and caching_sha2_password fast authentication.
struct MySQLWireError: LocalizedError {
    let message: String
    var errorDescription: String? { message }
}

private struct ByteCursor {
    let bytes: [UInt8]
    var offset = 0
    mutating func take(_ n: Int) throws -> [UInt8] {
        guard n >= 0, offset + n <= bytes.count else { throw MySQLWireError(message: "Malformed MySQL packet") }
        defer { offset += n }
        return Array(bytes[offset..<(offset+n)])
    }
    mutating func u8() throws -> UInt8 { try take(1)[0] }
    mutating func u16() throws -> Int { let b = try take(2); return Int(b[0]) | (Int(b[1]) << 8) }
    mutating func u32() throws -> Int { let b = try take(4); return Int(b[0]) | (Int(b[1]) << 8) | (Int(b[2]) << 16) | (Int(b[3]) << 24) }
    mutating func nulString() throws -> String {
        let start = offset
        while offset < bytes.count && bytes[offset] != 0 { offset += 1 }
        guard offset < bytes.count else { throw MySQLWireError(message: "Unterminated server string") }
        let s = String(decoding: bytes[start..<offset], as: UTF8.self)
        offset += 1
        return s
    }
    mutating func lenInt() throws -> Int? {
        let first = try u8()
        switch first {
        case 0xfb: return nil
        case 0xfc: return try u16()
        case 0xfd:
            let b = try take(3); return Int(b[0]) | Int(b[1]) << 8 | Int(b[2]) << 16
        case 0xfe:
            let b = try take(8)
            var v: UInt64 = 0
            for (i, byte) in b.enumerated() { v |= UInt64(byte) << (8 * i) }
            guard v <= UInt64(Int.max) else { throw MySQLWireError(message: "Field too large") }
            return Int(v)
        default: return Int(first)
        }
    }
    mutating func lenString() throws -> String? {
        guard let n = try lenInt() else { return nil }
        return String(decoding: try take(n), as: UTF8.self)
    }
}

private func le32(_ v: Int) -> [UInt8] { (0..<4).map { UInt8(truncatingIfNeeded: v >> ($0 * 8)) } }
private func sha1(_ data: [UInt8]) -> [UInt8] {
    var result = [UInt8](repeating: 0, count: Int(CC_SHA1_DIGEST_LENGTH))
    data.withUnsafeBytes { ptr in _ = CC_SHA1(ptr.baseAddress, CC_LONG(data.count), &result) }
    return result
}
private func sha256(_ data: [UInt8]) -> [UInt8] {
    var result = [UInt8](repeating: 0, count: Int(CC_SHA256_DIGEST_LENGTH))
    data.withUnsafeBytes { ptr in _ = CC_SHA256(ptr.baseAddress, CC_LONG(data.count), &result) }
    return result
}
private func authResponse(password: String, salt: [UInt8], plugin: String) throws -> [UInt8] {
    if password.isEmpty { return [] }
    let pwd = Array(password.utf8)
    if plugin == "mysql_native_password" {
        let first = sha1(pwd), second = sha1(sha1(pwd))
        let last = sha1(salt + second)
        return zip(first, last).map { $0 ^ $1 }
    }
    if plugin == "caching_sha2_password" {
        let first = sha256(pwd), second = sha256(sha256(pwd))
        let last = sha256(second + salt)
        return zip(first, last).map { $0 ^ $1 }
    }
    throw MySQLWireError(message: "Authentication plugin \(plugin) is not yet supported")
}

// Parse an RSA SubjectPublicKeyInfo PEM supplied by MySQL. Security.framework
// expects the inner PKCS#1 RSAPublicKey DER for SecKeyCreateWithData.
private func mysqlRSAPublicKey(_ packet: [UInt8]) throws -> SecKey {
    guard packet.first == 0x01 else {
        throw MySQLWireError(message: "Expected caching_sha2_password public-key response")
    }
    let pem = String(decoding: packet.dropFirst(), as: UTF8.self)
    let base64 = pem.components(separatedBy: .newlines)
        .filter { !$0.hasPrefix("-----") }.joined()
        .trimmingCharacters(in: .whitespacesAndNewlines)
    guard let spki = Data(base64Encoded: base64) else {
        throw MySQLWireError(message: "Invalid MySQL RSA public key")
    }
    let der = [UInt8](spki)
    func length(_ index: inout Int) throws -> Int {
        guard index < der.count else { throw MySQLWireError(message: "Invalid RSA DER length") }
        let first = Int(der[index]); index += 1
        if first < 128 { return first }
        let bytes = first & 0x7f
        guard bytes > 0, bytes <= 4, index + bytes <= der.count else {
            throw MySQLWireError(message: "Invalid RSA DER length")
        }
        var value = 0
        for _ in 0..<bytes { value = (value << 8) | Int(der[index]); index += 1 }
        return value
    }
    func element(_ index: inout Int, tag: UInt8) throws -> Range<Int> {
        guard index < der.count, der[index] == tag else {
            throw MySQLWireError(message: "Invalid RSA DER element")
        }
        index += 1
        let count = try length(&index)
        guard count <= der.count - index else { throw MySQLWireError(message: "Truncated RSA DER") }
        let range = index..<(index + count)
        index += count
        return range
    }
    var i = 0
    let outer = try element(&i, tag: 0x30)
    i = outer.lowerBound
    _ = try element(&i, tag: 0x30) // AlgorithmIdentifier
    let bits = try element(&i, tag: 0x03)
    guard bits.count > 1, der[bits.lowerBound] == 0 else {
        throw MySQLWireError(message: "Invalid RSA public key bit string")
    }
    let pkcs1 = Data(der[(bits.lowerBound + 1)..<bits.upperBound])
    var error: Unmanaged<CFError>?
    guard let key = SecKeyCreateWithData(pkcs1 as CFData, [
        kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
        kSecAttrKeyClass as String: kSecAttrKeyClassPublic
    ] as CFDictionary, &error) else {
        throw MySQLWireError(message: "Cannot import MySQL RSA public key: \(error?.takeRetainedValue().localizedDescription ?? "unknown error")")
    }
    return key
}

final class MySQLWireConnection {
    private let input: InputStream
    private let output: OutputStream
    private var sequence: UInt8 = 0
    private let maxPacket = 16_777_215

    init(host: String, port: Int) throws {
        guard (1...65535).contains(port) else { throw MySQLWireError(message: "Invalid port") }
        var read: Unmanaged<CFReadStream>?
        var write: Unmanaged<CFWriteStream>?
        CFStreamCreatePairWithSocketToHost(nil, host as CFString, UInt32(port), &read, &write)
        guard let r = read?.takeRetainedValue(), let w = write?.takeRetainedValue() else {
            throw MySQLWireError(message: "Cannot create connection")
        }
        input = unsafeBitCast(r, to: InputStream.self)
        output = unsafeBitCast(w, to: OutputStream.self)
        input.open(); output.open()
        // A CFStream may still be .opening immediately after open(). The
        // first blocking read below waits for the server greeting.
        if input.streamStatus == .error || output.streamStatus == .error {
            throw MySQLWireError(message: input.streamError?.localizedDescription ?? output.streamError?.localizedDescription ?? "Could not open MySQL connection")
        }
    }
    deinit { input.close(); output.close() }

    private func readExact(_ count: Int) throws -> [UInt8] {
        var buffer = [UInt8](repeating: 0, count: count)
        var done = 0
        while done < count {
            let n = buffer.withUnsafeMutableBufferPointer { ptr -> Int in
                guard let base = ptr.baseAddress else { return -1 }
                return input.read(base.advanced(by: done), maxLength: count - done)
            }
            guard n > 0 else { throw MySQLWireError(message: input.streamError?.localizedDescription ?? "MySQL connection closed") }
            done += n
        }
        return buffer
    }
    private func readPacket() throws -> [UInt8] {
        let header = try readExact(4)
        let count = Int(header[0]) | Int(header[1]) << 8 | Int(header[2]) << 16
        sequence = header[3] &+ 1
        return try readExact(count)
    }
    private func send(_ payload: [UInt8]) throws {
        guard payload.count < maxPacket else { throw MySQLWireError(message: "Query exceeds initial client packet limit") }
        let count = payload.count
        let packet = [UInt8(count & 255), UInt8((count >> 8) & 255), UInt8((count >> 16) & 255), sequence] + payload
        sequence = sequence &+ 1
        var written = 0
        while written < packet.count {
            let n = packet.withUnsafeBufferPointer { ptr -> Int in
                guard let base = ptr.baseAddress else { return -1 }
                return output.write(base.advanced(by: written), maxLength: packet.count - written)
            }
            guard n > 0 else { throw MySQLWireError(message: output.streamError?.localizedDescription ?? "MySQL write failed") }
            written += n
        }
    }
    private func checkError(_ p: [UInt8]) throws {
        guard p.first == 0xff else { return }
        let details = String(decoding: p.dropFirst(3), as: UTF8.self)
        throw MySQLWireError(message: "MySQL error: \(details)")
    }
    func authenticate(user: String, password: String, database: String) throws {
        let hello = try readPacket()
        var c = ByteCursor(bytes: hello)
        let version = try c.u8()
        guard version == 10 else { throw MySQLWireError(message: "Unsupported MySQL handshake") }
        _ = try c.nulString()
        _ = try c.u32()
        let salt1 = try c.take(8)
        _ = try c.u8()
        let lower = try c.u16()
        _ = try c.u8() // character set
        _ = try c.u16() // status
        let upper = try c.u16()
        let capabilities = lower | (upper << 16)
        let authLen = Int(try c.u8())
        _ = try c.take(10)
        let salt2 = try c.take(max(13, authLen - 8))
        let salt = salt1 + salt2.prefix(max(0, authLen - 9))
        let plugin = (try? c.nulString()) ?? "mysql_native_password"
        // This initial implementation deliberately refuses cleartext remote credentials.
        guard hostIsLoopback else { throw MySQLWireError(message: "Direct remote connections require TLS, which is not yet implemented. Use a local SSH tunnel.") }
        // CLIENT_LONG_PASSWORD, LONG_FLAG, PROTOCOL_41, TRANSACTIONS,
        // SECURE_CONNECTION and PLUGIN_AUTH. Do not set DEPRECATE_EOF.
        let clientFlags = 0x00000001 | 0x00000004 | 0x00000200 | 0x00002000 | 0x00008000 | 0x00080000
        let response = try authResponse(password: password, salt: Array(salt), plugin: plugin)
        var body = le32(clientFlags & capabilities) + le32(16_777_216) + [UInt8(45)] + [UInt8](repeating: 0, count: 23)
        body += Array(user.utf8) + [0, UInt8(response.count)] + response
        body += Array(plugin.utf8) + [0]
        try send(body)
        var reply = try readPacket()
        try checkError(reply)
        var currentPlugin = plugin
        var currentSalt = Array(salt)
        if reply.first == 0xfe {
            var switchCursor = ByteCursor(bytes: Array(reply.dropFirst()))
            let newPlugin = try switchCursor.nulString()
            let rawSalt = Array(reply.dropFirst(1 + newPlugin.utf8.count + 1))
            let newSalt = rawSalt.last == 0 ? Array(rawSalt.dropLast()) : rawSalt
            currentPlugin = newPlugin
            currentSalt = newSalt
            try send(authResponse(password: password, salt: newSalt, plugin: newPlugin))
            reply = try readPacket()
            try checkError(reply)
        }
        if reply.first == 0x01 {
            guard currentPlugin == "caching_sha2_password", reply.count >= 2 else {
                throw MySQLWireError(message: "Unsupported authentication continuation")
            }
            switch reply[1] {
            case 0x03: // Fast authentication accepted; wait for OK.
                reply = try readPacket()
            case 0x04: // Full authentication required.
                guard hostIsLoopback else {
                    throw MySQLWireError(message: "Remote full authentication requires TLS (not implemented). Use a trusted local SSH tunnel.")
                }
                if password.isEmpty {
                    try send([0])
                } else {
                    // Request the server's RSA public key (0x02). This is
                    // intended ONLY for a trusted local loopback connection.
                    try send([0x02])
                    let keyPacket = try readPacket()
                    try checkError(keyPacket)
                    let key = try mysqlRSAPublicKey(keyPacket)
                    guard !currentSalt.isEmpty else {
                        throw MySQLWireError(message: "Missing MySQL authentication nonce")
                    }
                    var scrambled = Array(password.utf8) + [0]
                    for i in scrambled.indices { scrambled[i] ^= currentSalt[i % currentSalt.count] }
                    var rsaError: Unmanaged<CFError>?
                    guard let encrypted = SecKeyCreateEncryptedData(
                        key, .rsaEncryptionOAEPSHA1, Data(scrambled) as CFData, &rsaError
                    ) as Data? else {
                        throw MySQLWireError(message: "MySQL RSA encryption failed: \(rsaError?.takeRetainedValue().localizedDescription ?? "unknown error")")
                    }
                    try send(Array(encrypted))
                }
                reply = try readPacket()
            default:
                throw MySQLWireError(message: "Unsupported caching_sha2_password response: \(reply[1])")
            }
        }
        try checkError(reply)
        guard reply.first == 0x00 else { throw MySQLWireError(message: "Authentication did not complete") }
        if !database.isEmpty { _ = try query("USE `\(database.replacingOccurrences(of: "`", with: "``"))`") }
    }
    private var hostIsLoopback = false
    func markHost(_ host: String) { hostIsLoopback = ["localhost", "127.0.0.1", "::1"].contains(host.lowercased()) }

    func query(_ sql: String) throws -> QueryResult {
        sequence = 0
        try send([0x03] + Array(sql.utf8))
        let first = try readPacket()
        try checkError(first)
        if first.first == 0x00 { return QueryResult(message: "Statement executed successfully") }
        var c = ByteCursor(bytes: first)
        guard let columnCount = try c.lenInt(), columnCount <= 4096 else { throw MySQLWireError(message: "Invalid column count") }
        var names: [String] = []
        for _ in 0..<columnCount {
            var field = ByteCursor(bytes: try readPacket())
            for _ in 0..<4 { _ = try field.lenString() }
            names.append(try field.lenString() ?? "")
        }
        _ = try readPacket() // column EOF (initial protocol requires EOF)
        var rows: [[String]] = []
        while true {
            let packet = try readPacket()
            try checkError(packet)
            if packet.first == 0xfe && packet.count < 9 { break }
            var row = ByteCursor(bytes: packet)
            var values: [String] = []
            for _ in 0..<columnCount { values.append(try row.lenString() ?? "NULL") }
            rows.append(values)
            if rows.count > 100_000 { throw MySQLWireError(message: "Result exceeds 100,000 row safety limit") }
        }
        return QueryResult(columns: names, rows: rows, message: "\(rows.count) row(s)")
    }
}


// A session owns one authenticated TCP connection. Access is serialized by the
// queue below because MySQLWireConnection is not safe for concurrent queries.
private final class PersistentMySQLSessions {
    static let shared = PersistentMySQLSessions()

    private struct Entry {
        let profile: ConnectionProfile
        let password: String
        let connection: MySQLWireConnection
    }

    private let queue = DispatchQueue(label: "WorkbenchReplacement.MySQLSessions")
    private var connections: [UUID: Entry] = [:]

    func execute(profile: ConnectionProfile, database: String,
                 password: String, sql: String) throws -> QueryResult {
        try queue.sync {
            // A changed connection profile or password requires a new session.
            if let previous = connections[profile.id],
               (previous.profile != profile || previous.password != password) {
                connections.removeValue(forKey: profile.id)
            }

            if connections[profile.id] == nil {
                let connection = try MySQLWireConnection(host: profile.host, port: profile.port)
                connection.markHost(profile.host)
                try connection.authenticate(user: profile.username,
                                            password: password, database: database)
                connections[profile.id] = Entry(profile: profile,
                                                password: password,
                                                connection: connection)
            }

            guard let entry = connections[profile.id] else {
                throw MySQLWireError(message: "MySQL session was not created")
            }
            do {
                // Crucially, no new authentication takes place here. USE,
                // session variables and transactions survive between queries.
                return try entry.connection.query(sql)
            } catch {
                // The wire protocol can be out of sync after a failed query.
                // Discard it rather than silently issuing commands on a bad socket.
                connections.removeValue(forKey: profile.id)
                throw error
            }
        }
    }

    func disconnect(profileID: UUID) {
        queue.sync { _ = connections.removeValue(forKey: profileID) }
    }
}

enum NativeMySQLRunner {
    static func run(profile: ConnectionProfile, database: String,
                    password: String, sql: String) async throws -> QueryResult {
        try await Task.detached(priority: .userInitiated) {
            try PersistentMySQLSessions.shared.execute(
                profile: profile, database: database, password: password, sql: sql
            )
        }.value
    }

    static func disconnect(profileID: UUID) {
        PersistentMySQLSessions.shared.disconnect(profileID: profileID)
    }
}