import Foundation
import CommonCrypto
import Compression
import zlib

// MARK: - CryptoPrimitives

/// Shared low-level cryptographic primitives used by both `CryptoManager` and `MessageCrypto`.
/// Centralizes AES-CBC, PBKDF2, random bytes, and zlib to avoid duplication.
enum CryptoPrimitives {

    // MARK: - AES-256-CBC

    nonisolated static func aesCBCEncrypt(_ data: Data, key: Data, iv: Data) throws -> Data {
        let outputSize = data.count + kCCBlockSizeAES128
        var ciphertext = Data(count: outputSize)
        var numBytes = 0
        let status = ciphertext.withUnsafeMutableBytes { ciphertextPtr in
            data.withUnsafeBytes { dataPtr in
                key.withUnsafeBytes { keyPtr in
                    iv.withUnsafeBytes { ivPtr in
                        guard let cBase = ciphertextPtr.baseAddress,
                              let dBase = dataPtr.baseAddress,
                              let kBase = keyPtr.baseAddress,
                              let iBase = ivPtr.baseAddress else { return CCCryptorStatus(kCCParamError) }
                        return CCCrypt(
                            CCOperation(kCCEncrypt),
                            CCAlgorithm(kCCAlgorithmAES),
                            CCOptions(kCCOptionPKCS7Padding),
                            kBase, key.count,
                            iBase,
                            dBase, data.count,
                            cBase, outputSize,
                            &numBytes
                        )
                    }
                }
            }
        }
        guard status == kCCSuccess else { throw CryptoError.encryptionFailed }
        return ciphertext.prefix(numBytes)
    }

    nonisolated static func aesCBCDecrypt(_ data: Data, key: Data, iv: Data) throws -> Data {
        let outputSize = data.count + kCCBlockSizeAES128
        var plaintext = Data(count: outputSize)
        var numBytes = 0
        let status = plaintext.withUnsafeMutableBytes { plaintextPtr in
            data.withUnsafeBytes { dataPtr in
                key.withUnsafeBytes { keyPtr in
                    iv.withUnsafeBytes { ivPtr in
                        guard let pBase = plaintextPtr.baseAddress,
                              let dBase = dataPtr.baseAddress,
                              let kBase = keyPtr.baseAddress,
                              let iBase = ivPtr.baseAddress else { return CCCryptorStatus(kCCParamError) }
                        return CCCrypt(
                            CCOperation(kCCDecrypt),
                            CCAlgorithm(kCCAlgorithmAES),
                            CCOptions(kCCOptionPKCS7Padding),
                            kBase, key.count,
                            iBase,
                            dBase, data.count,
                            pBase, outputSize,
                            &numBytes
                        )
                    }
                }
            }
        }
        guard status == kCCSuccess else { throw CryptoError.decryptionFailed }
        return plaintext.prefix(numBytes)
    }

    // MARK: - PBKDF2

    nonisolated static func pbkdf2(
        password: String,
        salt: String,
        iterations: Int,
        keyLength: Int,
        prf: CCPseudoRandomAlgorithm
    ) -> Data {
        // Desktop parity: CryptoJS PBKDF2 encodes the full UTF-8 string
        // including embedded null bytes (U+0000). Swift's withCString +
        // strlen() truncates at the first null byte, producing a different
        // PBKDF2 key. Use explicit UTF-8 byte arrays to match JavaScript.
        let passwordBytes = Array(password.utf8)
        let saltBytes = Array(salt.utf8)
        var derivedKey = Data(repeating: 0, count: keyLength)
        derivedKey.withUnsafeMutableBytes { keyPtr in
            guard let keyBase = keyPtr.bindMemory(to: UInt8.self).baseAddress else { return }
            passwordBytes.withUnsafeBufferPointer { passPtr in
                saltBytes.withUnsafeBufferPointer { saltPtr in
                    _ = CCKeyDerivationPBKDF(
                        CCPBKDFAlgorithm(kCCPBKDF2),
                        passPtr.baseAddress, passPtr.count,
                        saltPtr.baseAddress, saltPtr.count,
                        prf,
                        UInt32(iterations),
                        keyBase,
                        keyLength
                    )
                }
            }
        }
        return derivedKey
    }

    // MARK: - Random Bytes

    nonisolated static func randomBytes(count: Int) throws -> Data {
        var data = Data(count: count)
        let status = data.withUnsafeMutableBytes { ptr -> OSStatus in
            guard let base = ptr.baseAddress else { return errSecAllocate }
            return SecRandomCopyBytes(kSecRandomDefault, count, base)
        }
        guard status == errSecSuccess else { throw CryptoError.invalidEntropy }
        return data
    }
}

// MARK: - zlib Deflate / Inflate

extension CryptoPrimitives {

    /// Zlib-wrapped deflate compression using the system C zlib library (`compress2()`).
    /// Produces standard RFC 1950 zlib output compatible with desktop `pako.inflate()`
    /// and Node.js `zlib.inflateSync()`.
    ///
    /// Previously used Apple's `compression_encode_buffer(COMPRESSION_ZLIB)` (raw deflate)
    /// with a manual zlib wrapper — that output was incompatible with pako.inflate().
    nonisolated static func zlibDeflate(_ data: Data) throws -> Data {
        let sourceLen = uLong(data.count)
        var destLen = compressBound(sourceLen)
        var dest = Data(count: Int(destLen))

        let status = dest.withUnsafeMutableBytes { destPtr in
            data.withUnsafeBytes { srcPtr in
                guard let dBase = destPtr.bindMemory(to: Bytef.self).baseAddress,
                      let sBase = srcPtr.bindMemory(to: Bytef.self).baseAddress else {
                    return Z_MEM_ERROR
                }
                return compress2(dBase, &destLen, sBase, sourceLen, Z_DEFAULT_COMPRESSION)
            }
        }
        guard status == Z_OK else { throw CryptoError.compressionFailed }
        return dest.prefix(Int(destLen))
    }

    /// Raw deflate compression (no zlib header, compatible with Java Deflater(nowrap=true)).
    nonisolated static func rawDeflate(_ data: Data) throws -> Data {
        let sourceSize = data.count
        let destinationSize = sourceSize + 512
        var destination = Data(count: destinationSize)

        let compressedSize = destination.withUnsafeMutableBytes { destPtr in
            data.withUnsafeBytes { srcPtr in
                guard let dBase = destPtr.bindMemory(to: UInt8.self).baseAddress,
                      let sBase = srcPtr.bindMemory(to: UInt8.self).baseAddress else { return 0 }
                return compression_encode_buffer(dBase, destinationSize, sBase, sourceSize, nil, COMPRESSION_ZLIB)
            }
        }
        guard compressedSize > 0 else { throw CryptoError.compressionFailed }
        return destination.prefix(compressedSize)
    }

    /// Inflate decompression supporting both raw deflate and zlib-wrapped formats.
    /// Apple's `COMPRESSION_ZLIB` uses windowBits=-15 (raw deflate only).
    /// Desktop `pako.deflate()` / iOS `zlibDeflate()` produce zlib-wrapped data
    /// (0x78 header + raw deflate + adler32 trailer).
    ///
    /// ⚠️ zlib-wrapped data MUST be stripped FIRST — `tryRawInflate` can produce
    /// garbage output from zlib header bytes (0x78 0x9C are valid but meaningless
    /// raw deflate instructions), causing false-positive decompression.
    nonisolated static func rawInflate(_ data: Data) throws -> Data {
        // 1. Strip zlib wrapper FIRST if present (iOS zlibDeflate / desktop pako.deflate).
        //    Must be tried before raw inflate to avoid false-positive decompression
        //    where raw inflate interprets the zlib header as deflate instructions.
        if data.count > 6, data[data.startIndex] == 0x78 {
            let stripped = Data(data[(data.startIndex + 2) ..< (data.endIndex - 4)])
            if let result = tryRawInflate(stripped) {
                return result
            }
        }
        // 2. Try raw deflate (legacy iOS rawDeflate / Android-produced)
        if let result = tryRawInflate(data) {
            return result
        }
        throw CryptoError.compressionFailed
    }

    nonisolated private static func tryRawInflate(_ data: Data) -> Data? {
        let sourceSize = data.count
        for multiplier in [4, 8, 16, 32] {
            let destinationSize = max(sourceSize * multiplier, 256)
            var destination = Data(count: destinationSize)
            let decompressedSize = destination.withUnsafeMutableBytes { destPtr in
                data.withUnsafeBytes { srcPtr in
                    guard let dBase = destPtr.bindMemory(to: UInt8.self).baseAddress,
                          let sBase = srcPtr.bindMemory(to: UInt8.self).baseAddress else { return 0 }
                    return compression_decode_buffer(dBase, destinationSize, sBase, sourceSize, nil, COMPRESSION_ZLIB)
                }
            }
            if decompressedSize > 0, decompressedSize < destinationSize {
                return destination.prefix(decompressedSize)
            }
        }
        return nil
    }
}

// MARK: - Data Extensions

extension Data {
    /// Lowercase hex string representation.
    nonisolated var hexString: String {
        map { String(format: "%02x", $0) }.joined()
    }

    /// Initialize from a hex string (case-insensitive).
    nonisolated init(hexString: String) {
        let hex = hexString.lowercased()
        var data = Data(capacity: hex.count / 2)
        var index = hex.startIndex
        while index < hex.endIndex {
            let nextIndex = hex.index(index, offsetBy: 2, limitedBy: hex.endIndex) ?? hex.endIndex
            if nextIndex == hex.endIndex && hex.distance(from: index, to: nextIndex) < 2 {
                let byte = UInt8(hex[index...index], radix: 16) ?? 0
                data.append(byte)
            } else {
                let byte = UInt8(hex[index..<nextIndex], radix: 16) ?? 0
                data.append(byte)
            }
            index = nextIndex
        }
        self = data
    }

    /// Initialize from a STRICT hex string.
    /// Returns nil if the input contains non-hex characters or has odd length.
    nonisolated init?(strictHexString: String) {
        let hex = strictHexString.trimmingCharacters(in: .whitespacesAndNewlines)
        guard hex.count % 2 == 0 else { return nil }

        var data = Data(capacity: hex.count / 2)
        var index = hex.startIndex
        while index < hex.endIndex {
            let nextIndex = hex.index(index, offsetBy: 2)
            guard let byte = UInt8(hex[index..<nextIndex], radix: 16) else {
                return nil
            }
            data.append(byte)
            index = nextIndex
        }
        self = data
    }
}
