import Foundation
import CommonCrypto
import P256K

// MARK: - MessageCrypto

/// Handles message-level encryption/decryption using XChaCha20-Poly1305 + ECDH.
/// Matches the Android `MessageCrypto` implementation for cross-platform compatibility.
///
/// Crypto engines are split into dedicated files:
/// - `XChaCha20Engine` — ChaCha20 stream cipher + XChaCha20-Poly1305 AEAD
/// - `Poly1305Engine` — Poly1305 MAC
/// - `CryptoPrimitives` — Shared AES-CBC, PBKDF2, zlib, hex helpers
enum MessageCrypto {

    // MARK: - Public API

    /// Decrypts an incoming message using ECDH key exchange + XChaCha20-Poly1305.
    /// - Parameters:
    ///   - ciphertext: Hex-encoded XChaCha20-Poly1305 encrypted content (ciphertext + 16-byte tag).
    ///   - encryptedKey: Base64-encoded ECDH-encrypted key+nonce (`iv:encryptedKey:ephemeralPrivateKey`).
    ///   - myPrivateKeyHex: Recipient's secp256k1 private key (hex).
    /// - Returns: Decrypted plaintext string.
    /// Decrypts an incoming message and returns both plaintext and the working key+nonce.
    /// The returned `keyAndNonce` is the candidate that successfully decrypted the message —
    /// critical for deriving the correct attachment password.
    nonisolated static func decryptIncomingFull(
        ciphertext: String,
        encryptedKey: String,
        myPrivateKeyHex: String
    ) throws -> (text: String, keyAndNonce: Data) {
        let keyCandidates = try decryptKeyFromSenderCandidates(
            encryptedKey: encryptedKey,
            myPrivateKeyHex: myPrivateKeyHex
        )

        var lastError: Error?
        for keyAndNonce in keyCandidates where keyAndNonce.count >= 56 {
            do {
                let text = try decryptWithKeyAndNonce(ciphertext: ciphertext, keyAndNonce: keyAndNonce)
                return (text, keyAndNonce)
            } catch {
                lastError = error
            }
        }

        if let lastError {
            throw lastError
        }
        throw CryptoError.invalidData("Failed to decrypt message content with all key candidates")
    }

    nonisolated static func decryptIncoming(
        ciphertext: String,
        encryptedKey: String,
        myPrivateKeyHex: String
    ) throws -> String {
        try decryptIncomingFull(
            ciphertext: ciphertext,
            encryptedKey: encryptedKey,
            myPrivateKeyHex: myPrivateKeyHex
        ).text
    }

    /// Encrypts a message using XChaCha20-Poly1305 + ECDH for the recipient.
    /// - Parameters:
    ///   - plaintext: The message text.
    ///   - recipientPublicKeyHex: Recipient's secp256k1 compressed public key (hex).
    /// - Returns: Tuple of (content: hex ciphertext+tag, chachaKey: base64 encrypted key for recipient, plainKeyAndNonce: raw key+nonce bytes).
    nonisolated static func encryptOutgoing(
        plaintext: String,
        recipientPublicKeyHex: String
    ) throws -> (content: String, chachaKey: String, plainKeyAndNonce: Data) {
        guard let plaintextData = plaintext.data(using: .utf8) else {
            throw CryptoError.invalidData("Cannot encode plaintext as UTF-8")
        }

        let key = try CryptoPrimitives.randomBytes(count: 32)
        let nonce = try CryptoPrimitives.randomBytes(count: 24)
        let keyAndNonce = key + nonce

        let ciphertextWithTag = try XChaCha20Engine.encrypt(
            plaintext: plaintextData, key: key, nonce: nonce
        )

        let chachaKey = try encryptKeyForRecipient(
            keyAndNonce: keyAndNonce, recipientPublicKeyHex: recipientPublicKeyHex
        )

        return (
            content: ciphertextWithTag.hexString,
            chachaKey: chachaKey,
            plainKeyAndNonce: keyAndNonce
        )
    }

    /// Decrypts an incoming message using already decrypted key+nonce bytes.
    /// Mirrors Android `decryptIncomingWithPlainKey`.
    nonisolated static func decryptIncomingWithPlainKey(
        ciphertext: String,
        plainKeyAndNonce: Data
    ) throws -> String {
        guard plainKeyAndNonce.count >= 56 else {
            throw CryptoError.invalidData("Key+nonce must be 56 bytes, got \(plainKeyAndNonce.count)")
        }
        return try decryptWithKeyAndNonce(ciphertext: ciphertext, keyAndNonce: plainKeyAndNonce)
    }

    /// Extract raw decrypted key+nonce data from an encrypted key string (ECDH path).
    /// Verifies each candidate by attempting XChaCha20 decryption to find the correct one.
    /// Falls back to first candidate if ciphertext is unavailable.
    nonisolated static func extractDecryptedKeyData(
        encryptedKey: String,
        myPrivateKeyHex: String,
        verifyCiphertext: String? = nil
    ) -> Data? {
        guard let candidates = try? decryptKeyFromSenderCandidates(
            encryptedKey: encryptedKey, myPrivateKeyHex: myPrivateKeyHex
        ) else { return nil }
        // When ciphertext is available, verify each candidate via XChaCha20 decryption
        if let ciphertext = verifyCiphertext {
            for candidate in candidates where candidate.count >= 56 {
                if (try? decryptWithKeyAndNonce(ciphertext: ciphertext, keyAndNonce: candidate)) != nil {
                    return candidate
                }
            }
        }
        return candidates.first
    }

    /// Android parity helper:
    /// `String(bytes, UTF_8)` (replacement for invalid sequences) + `toByteArray(ISO_8859_1)`.
    /// Emulates Android's `String(bytes, UTF_8).toByteArray(ISO_8859_1)` round-trip.
    /// Uses BOTH WHATWG and Android UTF-8 decoders — returns candidates for each.
    /// WHATWG and Android decoders handle invalid UTF-8 differently → different bytes.
    nonisolated static func androidUtf8BytesToLatin1Bytes(_ utf8Bytes: Data) -> Data {
        // Primary: WHATWG decoder (matches Java's Modified UTF-8 for most cases)
        let decoded = String(decoding: utf8Bytes, as: UTF8.self)
        if let latin1 = decoded.data(using: .isoLatin1, allowLossyConversion: true) {
            return latin1
        }
        return Data(decoded.unicodeScalars.map { $0.value <= 0xFF ? UInt8($0.value) : UInt8(ascii: "?") })
    }

    /// Alternative key recovery using Android UTF-8 decoder.
    nonisolated static func androidUtf8BytesToLatin1BytesAlt(_ utf8Bytes: Data) -> Data {
        let decoded = bytesToAndroidUtf8String(utf8Bytes)
        if let latin1 = decoded.data(using: .isoLatin1, allowLossyConversion: true) {
            return latin1
        }
        return Data(decoded.unicodeScalars.map { $0.value <= 0xFF ? UInt8($0.value) : UInt8(ascii: "?") })
    }

    // MARK: - Attachment Password Helpers

    /// Returns password candidates from a stored attachment password string.
    /// New format: `"rawkey:<hex>"` → derives HEX (primary) + Android/WHATWG/Latin-1 (backward compat).
    /// Legacy format: plain string → used as-is (backward compat with persisted messages).
    ///
    /// Desktop commit `61e83bd`: attachment password changed from `Buffer.toString('utf-8')` to
    /// `key.toString('hex')`. HEX is now the primary candidate; UTF-8 variants kept for old messages.
    nonisolated static func attachmentPasswordCandidates(from stored: String) -> [String] {
        if stored.hasPrefix("rawkey:") {
            let hex = String(stored.dropFirst("rawkey:".count))
            let keyData = Data(hexString: hex)
            let hexPwd = hex  // HEX: primary candidate (desktop commit 61e83bd)
            let androidPwd = bytesToAndroidUtf8String(keyData)
            let whatwgPwd = String(decoding: keyData, as: UTF8.self)
            // Latin-1 variant: backward compat with iOS builds that used .isoLatin1 encoding
            let latin1Pwd = String(bytes: keyData, encoding: .isoLatin1)
            var candidates = [hexPwd, androidPwd, whatwgPwd]
            if let latin1Pwd, latin1Pwd != hexPwd, latin1Pwd != androidPwd, latin1Pwd != whatwgPwd {
                candidates.append(latin1Pwd)
            }
            // Deduplicate while preserving order
            var seen = Set<String>()
            return candidates.filter { seen.insert($0).inserted }
        }
        // Group key or plain password — Android/iOS use plain groupKey.
        // Desktop SENDS with hex but RECEIVES with plain (Desktop bug).
        // Candidates include both hex and plain for backward compat.
        // If stored is already hex-encoded (128+ chars, all hex digits), use as-is
        // to avoid generating a 256-char double-hex garbage candidate.
        let isAlreadyHex = stored.count >= 128 && stored.allSatisfy { $0.isHexDigit }
        if isAlreadyHex {
            // Try hex as-is (Desktop) + un-hexed original (iOS/Android plain group key)
            let decoded = Data(hexString: stored)
            if let original = String(data: decoded, encoding: .utf8), !original.isEmpty {
                return [stored, original]
            }
            return [stored]
        }
        let hexVariant = Data(stored.utf8).map { String(format: "%02x", $0) }.joined()
        return [hexVariant, stored]
    }

    // MARK: - Android-Compatible UTF-8 Decoder

    /// Port of Android `bytesToJsUtf8String()` — manual UTF-8 decoder that emits
    /// ONE U+FFFD PER BYTE consumed in failed multi-byte sequences.
    ///
    /// This differs from WHATWG (Swift's `String(decoding:as:UTF8.self)`) which emits
    /// ONE U+FFFD per "maximal subpart" of an ill-formed subsequence.
    ///
    /// Example: bytes `[0xEF, 0x98, 0xF7]` (3-byte starter, valid continuation, invalid 2nd):
    /// - WHATWG: 1× U+FFFD (for maximal subpart [EF,98]), then re-scan F7
    /// - Android: 2× U+FFFD (one for EF, one for 98), then re-scan F7
    ///
    /// Used for PBKDF2 password derivation from random key+nonce bytes
    /// to match Android's attachment encryption.
    /// Exact port of Android's `bytesToBufferPolyfillUtf8String()` from MessageCrypto.kt:1419-1502.
    /// Uses feross/buffer npm polyfill UTF-8 decoding semantics.
    /// Key difference from previous implementation: on invalid sequence, ALWAYS advance by 1 byte
    /// and emit 1× U+FFFD (not variable bytes/U+FFFD count).
    nonisolated static func bytesToAndroidUtf8String(_ bytes: Data) -> String {
        var codePoints: [Int] = []
        codePoints.reserveCapacity(bytes.count)
        var index = 0
        let end = bytes.count

        while index < end {
            let firstByte = Int(bytes[index]) & 0xFF
            var codePoint: Int? = nil
            var bytesPerSequence: Int

            if firstByte > 0xEF {
                bytesPerSequence = 4
            } else if firstByte > 0xDF {
                bytesPerSequence = 3
            } else if firstByte > 0xBF {
                bytesPerSequence = 2
            } else {
                bytesPerSequence = 1
            }

            if index + bytesPerSequence <= end {
                switch bytesPerSequence {
                case 1:
                    if firstByte < 0x80 {
                        codePoint = firstByte
                    }
                case 2:
                    let secondByte = Int(bytes[index + 1]) & 0xFF
                    if (secondByte & 0xC0) == 0x80 {
                        let temp = ((firstByte & 0x1F) << 6) | (secondByte & 0x3F)
                        if temp > 0x7F {
                            codePoint = temp
                        }
                    }
                case 3:
                    let secondByte = Int(bytes[index + 1]) & 0xFF
                    let thirdByte = Int(bytes[index + 2]) & 0xFF
                    if (secondByte & 0xC0) == 0x80 && (thirdByte & 0xC0) == 0x80 {
                        let temp = ((firstByte & 0x0F) << 12) | ((secondByte & 0x3F) << 6) | (thirdByte & 0x3F)
                        if temp > 0x7FF && (temp < 0xD800 || temp > 0xDFFF) {
                            codePoint = temp
                        }
                    }
                case 4:
                    let secondByte = Int(bytes[index + 1]) & 0xFF
                    let thirdByte = Int(bytes[index + 2]) & 0xFF
                    let fourthByte = Int(bytes[index + 3]) & 0xFF
                    if (secondByte & 0xC0) == 0x80 && (thirdByte & 0xC0) == 0x80 && (fourthByte & 0xC0) == 0x80 {
                        let temp = ((firstByte & 0x0F) << 18) | ((secondByte & 0x3F) << 12) | ((thirdByte & 0x3F) << 6) | (fourthByte & 0x3F)
                        if temp > 0xFFFF && temp < 0x110000 {
                            codePoint = temp
                        }
                    }
                default:
                    break
                }
            }

            if codePoint == nil {
                codePoint = 0xFFFD
                bytesPerSequence = 1
            }

            codePoints.append(codePoint!)
            index += bytesPerSequence
        }

        // Build string from code points. Swift uses Unicode scalars (not UTF-16),
        // so supplementary plane characters (> 0xFFFF) are handled directly —
        // no surrogate pair decomposition needed (unlike JS/Kotlin UTF-16 strings).
        return String(codePoints.map { Character(UnicodeScalar($0)!) })
    }
}

// MARK: - ECDH Key Exchange

extension MessageCrypto {

    /// Decrypts and returns candidate XChaCha20 key+nonce buffers.
    /// Format: Base64(ivHex:encryptedKeyHex:ephemeralPrivateKeyHex)
    /// Supports Android sync shorthand `sync:<aesChachaKey>`.
    nonisolated static func decryptKeyFromSenderCandidates(encryptedKey: String, myPrivateKeyHex: String) throws -> [Data] {
        if encryptedKey.hasPrefix("sync:") {
            let aesChachaKey = String(encryptedKey.dropFirst("sync:".count))
            guard !aesChachaKey.isEmpty else {
                throw CryptoError.invalidData("Invalid sync key format")
            }
            let decrypted = try CryptoManager.shared.decryptWithPassword(
                aesChachaKey,
                password: myPrivateKeyHex
            )
            return [androidUtf8BytesToLatin1Bytes(decrypted)]
        }

        guard let decoded = Data(base64Encoded: encryptedKey) else {
            throw CryptoError.invalidData("Cannot decode Base64 key")
        }
        let combined = String(decoding: decoded, as: UTF8.self)

        let parts = combined.split(separator: ":", maxSplits: 2).map(String.init)
        guard parts.count == 3 else {
            throw CryptoError.invalidData("Expected iv:encrypted:ephemeralKey, got \(parts.count) parts")
        }

        let ivHex = parts[0]
        let encryptedKeyHex = parts[1]
        var ephemeralPrivateKeyHex = parts[2]

        // Handle odd-length hex from JS toString(16)
        if ephemeralPrivateKeyHex.count % 2 != 0 {
            ephemeralPrivateKeyHex = "0" + ephemeralPrivateKeyHex
        }

        let iv = Data(hexString: ivHex)
        let encryptedKeyData = Data(hexString: encryptedKeyHex)

        let normalizedEphemeralPrivateKeyHex = normalizePrivateKeyHex(ephemeralPrivateKeyHex)
        let normalizedMyPrivateKeyHex = normalizePrivateKeyHex(myPrivateKeyHex)

        let ephemeralPrivKey = try P256K.KeyAgreement.PrivateKey(
            dataRepresentation: Data(hexString: normalizedEphemeralPrivateKeyHex), format: .compressed
        )
        let myPrivKey = try P256K.KeyAgreement.PrivateKey(
            dataRepresentation: Data(hexString: normalizedMyPrivateKeyHex), format: .compressed
        )

        // ECDH: ephemeralPrivateKey × myPublicKey → shared point
        let sharedSecret = try ephemeralPrivKey.sharedSecretFromKeyAgreement(
            with: myPrivKey.publicKey, format: .compressed
        )
        let sharedSecretData = sharedSecret.withUnsafeBytes { Data($0) }
        let sharedKeyExact = extractXCoordinate(from: sharedSecretData)
        let sharedKeyLegacy = legacySharedKey(fromExactX: sharedKeyExact)

        let candidateSharedKeys: [Data] = sharedKeyLegacy == sharedKeyExact
            ? [sharedKeyExact]
            : [sharedKeyExact, sharedKeyLegacy]

        var candidates: [Data] = []
        var seen: Set<Data> = []

        for sharedKey in candidateSharedKeys {
            guard let decryptedBytes = try? CryptoPrimitives.aesCBCDecrypt(
                encryptedKeyData,
                key: sharedKey,
                iv: iv
            ) else {
                continue
            }

            // Android parity: String(bytes, UTF_8) + toByteArray(ISO_8859_1)
            // Try BOTH WHATWG and Android UTF-8 decoders — they handle invalid
            // sequences differently, producing different raw key bytes.
            let originalBytes = androidUtf8BytesToLatin1Bytes(decryptedBytes)
            if originalBytes.count >= 56, seen.insert(originalBytes).inserted {
                candidates.append(originalBytes)
            }
            let altBytes = androidUtf8BytesToLatin1BytesAlt(decryptedBytes)
            if altBytes.count >= 56, altBytes != originalBytes, seen.insert(altBytes).inserted {
                candidates.append(altBytes)
            }
        }

        guard !candidates.isEmpty else {
            throw CryptoError.invalidData("Failed to decrypt key with all shared secret variants")
        }
        return candidates
    }

    /// Encrypts the XChaCha20 key+nonce for a recipient using ECDH.
    nonisolated static func encryptKeyForRecipient(keyAndNonce: Data, recipientPublicKeyHex: String) throws -> String {
        let ephemeralPrivKey = try P256K.KeyAgreement.PrivateKey()

        let recipientPubKey = try P256K.KeyAgreement.PublicKey(
            dataRepresentation: Data(hexString: recipientPublicKeyHex), format: .compressed
        )

        let sharedSecret = try ephemeralPrivKey.sharedSecretFromKeyAgreement(
            with: recipientPubKey, format: .compressed
        )
        let sharedSecretData = sharedSecret.withUnsafeBytes { Data($0) }
        let sharedKey = extractXCoordinate(from: sharedSecretData)

        // Android parity: String(bytes, ISO_8859_1) -> UTF-8 bytes.
        guard let latin1String = String(data: keyAndNonce, encoding: .isoLatin1) else {
            throw CryptoError.encryptionFailed
        }
        let utf8Representation = latin1String
        guard let dataToEncrypt = utf8Representation.data(using: .utf8) else {
            throw CryptoError.encryptionFailed
        }

        let iv = try CryptoPrimitives.randomBytes(count: 16)
        let ciphertext = try CryptoPrimitives.aesCBCEncrypt(dataToEncrypt, key: sharedKey, iv: iv)
        let ephemeralPrivKeyHex = ephemeralPrivKey.rawRepresentation.hexString

        let combined = "\(iv.hexString):\(ciphertext.hexString):\(ephemeralPrivKeyHex)"
        guard let base64 = combined.data(using: .utf8)?.base64EncodedString() else {
            throw CryptoError.encryptionFailed
        }
        return base64
    }

    /// Extracts the 32-byte x-coordinate from ECDH shared secret bytes.
    nonisolated static func extractXCoordinate(from sharedSecretData: Data) -> Data {
        // Uncompressed point: 0x04 || X(32) || Y(32)
        if sharedSecretData.count == 65, sharedSecretData.first == 0x04 {
            return sharedSecretData[1..<33]
        }
        // Compressed point: 0x02/0x03 || X(32)
        if sharedSecretData.count == 33 {
            return sharedSecretData[1..<33]
        }
        // Raw X coordinate (32 bytes)
        if sharedSecretData.count == 32 {
            return sharedSecretData
        }
        // Defensive fallback: keep last 32 bytes if shape is unexpected.
        if sharedSecretData.count > 32 {
            return sharedSecretData.suffix(32)
        }
        // Left-pad short variants.
        return Data(repeating: 0, count: 32 - sharedSecretData.count) + sharedSecretData
    }

    /// Legacy Android compatibility: x-coordinate serialized through BigInteger loses leading zeros.
    nonisolated static func legacySharedKey(fromExactX exactX: Data) -> Data {
        let trimmed = exactX.drop(while: { $0 == 0 })
        return trimmed.isEmpty ? Data([0]) : Data(trimmed)
    }

    /// JS/Android compatibility: private key hex can arrive without leading zero bytes.
    nonisolated static func normalizePrivateKeyHex(_ rawHex: String) -> String {
        var hex = rawHex
        if hex.count % 2 != 0 {
            hex = "0" + hex
        }
        if hex.count < 64 {
            return String(repeating: "0", count: 64 - hex.count) + hex
        }
        if hex.count > 64 {
            return String(hex.suffix(64))
        }
        return hex
    }

    nonisolated static func decryptWithKeyAndNonce(ciphertext: String, keyAndNonce: Data) throws -> String {
        guard keyAndNonce.count >= 56 else {
            throw CryptoError.invalidData("Key+nonce must be 56 bytes, got \(keyAndNonce.count)")
        }

        let key = keyAndNonce[0..<32]
        let nonce = keyAndNonce[32..<56]
        let ciphertextData = Data(hexString: ciphertext)
        let plaintext = try XChaCha20Engine.decrypt(
            ciphertextWithTag: ciphertextData,
            key: Data(key),
            nonce: Data(nonce)
        )

        guard let text = String(data: plaintext, encoding: .utf8) else {
            throw CryptoError.invalidData("Decrypted data is not valid UTF-8")
        }
        return text
    }
}
