import CommonCrypto
import CryptoKit
import Foundation

// MARK: - BigUInt

struct BigUInt {
    private(set) var words: [UInt32]  // little-endian (index 0 = least significant)

    static let zero = BigUInt(words: [0])
    static let one  = BigUInt(words: [1])

    init(words: [UInt32]) {
        self.words = words
        normalize()
    }

    init(decimal string: String) {
        var result = BigUInt.zero
        let ten = BigUInt(words: [10])
        for ch in string {
            guard let d = ch.wholeNumberValue else { continue }
            result = result.multiplied(by: ten)
            result = result.adding(BigUInt(words: [UInt32(d)]))
        }
        self = result
    }

    private mutating func normalize() {
        while words.count > 1 && words.last == 0 { words.removeLast() }
    }

    var isZero: Bool { words.count == 1 && words[0] == 0 }
    var isOdd: Bool  { words[0] & 1 == 1 }

    var bitWidth: Int {
        let n = words.count
        guard n > 0 else { return 0 }
        let top = words[n - 1]
        if top == 0 { return max(0, (n - 1) * 32) }
        return (n - 1) * 32 + (32 - top.leadingZeroBitCount)
    }

    func bit(_ index: Int) -> Bool {
        let word = index / 32
        let bit  = index % 32
        guard word < words.count else { return false }
        return (words[word] >> bit) & 1 == 1
    }

    // MARK: Comparison

    static func < (lhs: BigUInt, rhs: BigUInt) -> Bool {
        if lhs.words.count != rhs.words.count {
            return lhs.words.count < rhs.words.count
        }
        for i in stride(from: lhs.words.count - 1, through: 0, by: -1) {
            if lhs.words[i] != rhs.words[i] { return lhs.words[i] < rhs.words[i] }
        }
        return false
    }

    static func == (lhs: BigUInt, rhs: BigUInt) -> Bool { lhs.words == rhs.words }

    // MARK: Addition

    func adding(_ other: BigUInt) -> BigUInt {
        let maxLen = max(words.count, other.words.count)
        var result = [UInt32](repeating: 0, count: maxLen + 1)
        var carry: UInt64 = 0
        for i in 0..<maxLen {
            let a: UInt64 = i < words.count ? UInt64(words[i]) : 0
            let b: UInt64 = i < other.words.count ? UInt64(other.words[i]) : 0
            let sum = a + b + carry
            result[i] = UInt32(sum & 0xFFFFFFFF)
            carry = sum >> 32
        }
        result[maxLen] = UInt32(carry)
        return BigUInt(words: result)
    }

    // MARK: Subtraction (self >= other assumed)

    func subtracting(_ other: BigUInt) -> BigUInt {
        var result = words
        var borrow: Int64 = 0
        for i in 0..<result.count {
            let b: Int64 = i < other.words.count ? Int64(other.words[i]) : 0
            let diff = Int64(result[i]) - b - borrow
            if diff < 0 {
                result[i] = UInt32(bitPattern: Int32(truncatingIfNeeded: diff &+ 0x1_0000_0000))
                borrow = 1
            } else {
                result[i] = UInt32(diff)
                borrow = 0
            }
        }
        return BigUInt(words: result)
    }

    // MARK: Multiplication

    func multiplied(by other: BigUInt) -> BigUInt {
        let n = words.count
        let m = other.words.count
        var result = [UInt32](repeating: 0, count: n + m)
        for i in 0..<n {
            var carry: UInt64 = 0
            for j in 0..<m {
                let prod = UInt64(words[i]) * UInt64(other.words[j]) + UInt64(result[i + j]) + carry
                result[i + j] = UInt32(prod & 0xFFFFFFFF)
                carry = prod >> 32
            }
            result[i + m] += UInt32(carry)
        }
        return BigUInt(words: result)
    }

    // MARK: Modulo (binary shift-subtract)

    func modulo(_ divisor: BigUInt) -> BigUInt {
        if divisor.isZero { return .zero }
        if self < divisor { return self }

        var remainder = BigUInt.zero
        let totalBits = bitWidth

        for i in stride(from: totalBits - 1, through: 0, by: -1) {
            remainder = remainder.shiftedLeft1()
            if bit(i) {
                remainder.words[0] |= 1
            }
            if !(remainder < divisor) {
                remainder = remainder.subtracting(divisor)
            }
        }
        return remainder
    }

    private func shiftedLeft1() -> BigUInt {
        var result = [UInt32](repeating: 0, count: words.count + 1)
        var carry: UInt32 = 0
        for i in 0..<words.count {
            let shifted = (UInt64(words[i]) << 1) | UInt64(carry)
            result[i] = UInt32(shifted & 0xFFFFFFFF)
            carry = UInt32(shifted >> 32)
        }
        result[words.count] = carry
        return BigUInt(words: result)
    }

    // MARK: ModPow — right-to-left binary

    func modPow(exp: BigUInt, mod: BigUInt) -> BigUInt {
        if mod == .one { return .zero }
        var result = BigUInt.one
        var base   = modulo(mod)
        var e      = exp

        while !e.isZero {
            if e.isOdd {
                result = result.multiplied(by: base).modulo(mod)
            }
            e = e.shiftedRight1()
            base = base.multiplied(by: base).modulo(mod)
        }
        return result
    }

    private func shiftedRight1() -> BigUInt {
        var result = words
        var borrow: UInt32 = 0
        for i in stride(from: result.count - 1, through: 0, by: -1) {
            let new = (result[i] >> 1) | (borrow << 31)
            borrow = result[i] & 1
            result[i] = new
        }
        return BigUInt(words: result)
    }

    // MARK: Decimal String (O(n) per digit via UInt64 carry)

    var decimalString: String {
        if isZero { return "0" }
        var digits: [Character] = []
        var remaining = words

        while !(remaining.count == 1 && remaining[0] == 0) {
            var rem: UInt64 = 0
            for i in stride(from: remaining.count - 1, through: 0, by: -1) {
                let cur = (rem << 32) | UInt64(remaining[i])
                remaining[i] = UInt32(cur / 10)
                rem = cur % 10
            }
            digits.append(Character(String(rem)))
            while remaining.count > 1 && remaining.last == 0 { remaining.removeLast() }
        }
        return String(digits.reversed())
    }
}

// MARK: - MibCrypto

enum MibCrypto {
    static let defaultKey = "8M3L9SBF1AC4FRE56788M3L9SBF1AC4FRE5678"

    // DH exponent A
    private static let dhA = BigUInt(decimal: "1563516802667282387226490351799736881442299778484610378722158765594241028592123324764949712696577")

    // DH modulus P
    private static let dhP = BigUInt(decimal: "2410312426921032588552076022197566074856950548502459942654116941958108831682612228890093858261341614673227141477904012196503648957050582631942730706805009223062734745341073406696246014589361659774041027169249453200378729434170325843778659198143763193776859869524088940195577346119843545301547043747207749969763750084308926339295559968882457872412993810129130294592999947926365264059284647209730384947211681434464714438488520940127459844288859336526896320919633919")

    // Pre-computed: 2^A mod P (sent to server in every DH key exchange)
    static let cmod = "2301533261465719294935473752300816828871347570489984010719302185643154930578030858802510266223676458248149736055131586081546441805822467879511077637292661430926459449504165353261883687521672133746481188567585259995381323120449147180627189549882525129015625540437713248150022924125681681822594907090716721855161928482983925605401801685462643917325004120294159320617793870605573541358655167992228254033496751384810855107557385690184061912872019448968632814193704749"

    // Derives session key: Base64(SHA256(smod^A mod P as decimal string))
    static func deriveSessionKey(_ smod: String) -> String {
        let smodInt = BigUInt(decimal: smod)
        let shared = smodInt.modPow(exp: dhA, mod: dhP)
        let secretStr = shared.decimalString
        let sha = SHA256.hash(data: Data(secretStr.utf8))
        return Data(sha).base64EncodedString()
    }

    // MARK: Blowfish/ECB/PKCS7 Encrypt

    static func encrypt(_ json: [String: Any], key: String) throws -> String {
        guard let jsonData = try? JSONSerialization.data(withJSONObject: json) else {
            throw MibError.encryptionFailed
        }
        let keyBytes = key.data(using: .isoLatin1) ?? Data(key.utf8)
        let encrypted = try blowfishProcess(data: jsonData, key: keyBytes, operation: CCOperation(kCCEncrypt))
        return encrypted.base64EncodedString()
    }

    static func decrypt(_ base64: String, key: String) throws -> [String: Any] {
        guard let cipherData = Data(base64Encoded: base64, options: .ignoreUnknownCharacters) else {
            throw MibError.decryptionFailed
        }
        let keyBytes = key.data(using: .isoLatin1) ?? Data(key.utf8)
        let plainData = try blowfishProcess(data: cipherData, key: keyBytes, operation: CCOperation(kCCDecrypt))
        guard let json = try? JSONSerialization.jsonObject(with: plainData) as? [String: Any] else {
            throw MibError.decryptionFailed
        }
        return json
    }

    private static func blowfishProcess(data: Data, key: Data, operation: CCOperation) throws -> Data {
        var outLen = 0
        let bufLen = data.count + kCCBlockSizeBlowfish
        var outBuf = [UInt8](repeating: 0, count: bufLen)

        let status = data.withUnsafeBytes { dataPtr in
            key.withUnsafeBytes { keyPtr in
                CCCrypt(
                    operation,
                    CCAlgorithm(kCCAlgorithmBlowfish),
                    CCOptions(kCCOptionECBMode | kCCOptionPKCS7Padding),
                    keyPtr.baseAddress, key.count,
                    nil,
                    dataPtr.baseAddress, data.count,
                    &outBuf, bufLen,
                    &outLen
                )
            }
        }
        guard status == kCCSuccess else {
            throw operation == CCOperation(kCCEncrypt) ? MibError.encryptionFailed : MibError.decryptionFailed
        }
        return Data(outBuf[0..<outLen])
    }

    // MARK: SHA-256 helpers

    static func hashPassword(_ password: String) -> String {
        sha256Upper(password)
    }

    static func sha256Upper(_ input: String) -> String {
        let hash = SHA256.hash(data: Data(input.utf8))
        return hash.map { String(format: "%02X", $0) }.joined()
    }

    // pgf03 = SHA256(clientSalt + SHA256(passwordHash + userSalt))
    static func computePgf03(passwordHash: String, userSalt: String, clientSalt: String) -> String {
        let inner = sha256Upper(passwordHash + userSalt)
        return sha256Upper(clientSalt + inner)
    }
}