// SPDX-License-Identifier: MIT
// Copyright (c) 2026 Thibault Ducray
//
// This file is part of MyPwdTool's open-source sync/encryption core — see
// LICENSE-SYNC-CRYPTO.md at the repo root and https://tducray.fr/mypwdtool/open-source/.
// The rest of this application is proprietary and NOT covered by this license.

import Foundation
import CryptoKit

// MARK: - SyncCryptoError

enum SyncCryptoError: LocalizedError {
    case invalidSGKLength
    case invalidCiphertextFormat
    case encryptionFailed
    case decryptionFailed

    var errorDescription: String? {
        switch self {
        case .invalidSGKLength:        return "SGK must be 32 bytes"
        case .invalidCiphertextFormat: return "Ciphertext has invalid format"
        case .encryptionFailed:        return "SGK encryption failed"
        case .decryptionFailed:        return "SGK decryption failed — wrong key or tampered data"
        }
    }
}

// MARK: - SyncCrypto

enum SyncCrypto {

    // MARK: - Encrypt

    /// Encrypts payload JSON under the SGK.
    ///
    /// AAD = canonical UTF-8 of:
    ///   `messageId|streamId|senderDeviceId|hex(SHA256(sortedInboxIds joined by comma))|1`
    ///
    /// Returns a base64url-encoded blob: `[0x01][nonce:12][ciphertext][tag:16]`
    static func encrypt(
        payload: Data,
        sgk: Data,
        messageId: String,
        streamId: String,
        senderDeviceId: String,
        recipientInboxIds: [String],
        senderCounter: Int
    ) throws -> String {
        guard sgk.count == 32 else { throw SyncCryptoError.invalidSGKLength }

        let key = SymmetricKey(data: sgk)
        let aad = buildAAD(
            messageId: messageId,
            streamId: streamId,
            senderDeviceId: senderDeviceId,
            recipientInboxIds: recipientInboxIds
        )

        do {
            let sealed = try AES.GCM.seal(payload, using: key, authenticating: aad)
            guard let combined = sealed.combined else { throw SyncCryptoError.encryptionFailed }
            // Prefix with version byte 0x01
            var blob = Data([0x01])
            blob.append(combined)
            return base64urlEncode(blob)
        } catch let e as SyncCryptoError {
            throw e
        } catch {
            throw SyncCryptoError.encryptionFailed
        }
    }

    // MARK: - Decrypt

    /// Decrypts a base64url blob produced by `encrypt(...)`. Same AAD construction.
    static func decrypt(
        ciphertext: String,
        sgk: Data,
        messageId: String,
        streamId: String,
        senderDeviceId: String,
        recipientInboxIds: [String],
        senderCounter: Int
    ) throws -> Data {
        guard sgk.count == 32 else { throw SyncCryptoError.invalidSGKLength }
        guard var blob = base64urlDecode(ciphertext) else {
            throw SyncCryptoError.invalidCiphertextFormat
        }
        // Minimum: 1 (version) + 12 (nonce) + 16 (tag) = 29 bytes
        guard blob.count >= 29, blob[0] == 0x01 else {
            throw SyncCryptoError.invalidCiphertextFormat
        }
        blob.removeFirst() // strip version byte

        let key = SymmetricKey(data: sgk)
        let aad = buildAAD(
            messageId: messageId,
            streamId: streamId,
            senderDeviceId: senderDeviceId,
            recipientInboxIds: recipientInboxIds
        )

        do {
            let sealedBox = try AES.GCM.SealedBox(combined: blob)
            return try AES.GCM.open(sealedBox, using: key, authenticating: aad)
        } catch let e as SyncCryptoError {
            throw e
        } catch {
            throw SyncCryptoError.decryptionFailed
        }
    }

    // MARK: - Key generation

    /// Generate a new random 32-byte SGK.
    static func generateSGK() -> Data {
        var key = Data(count: 32)
        _ = key.withUnsafeMutableBytes { SecRandomCopyBytes(kSecRandomDefault, 32, $0.baseAddress!) }
        return key
    }

    // MARK: - AAD construction

    /// Build AAD bytes:
    ///   `messageId|streamId|senderDeviceId|hex(SHA256(sorted_inboxes_joined_by_comma))|1`
    ///
    /// Note: senderCounter is intentionally omitted from AAD (v1 simplification —
    /// it lives in the plaintext SyncPayload instead). This avoids bootstrap ordering issues.
    private static func buildAAD(
        messageId: String,
        streamId: String,
        senderDeviceId: String,
        recipientInboxIds: [String]
    ) -> Data {
        let sortedInboxes = recipientInboxIds.sorted().joined(separator: ",")
        let inboxesData = Data(sortedInboxes.utf8)
        let inboxesHash = SHA256.hash(data: inboxesData)
        let inboxesHex = inboxesHash.map { String(format: "%02x", $0) }.joined()

        let aadString = "\(messageId)|\(streamId)|\(senderDeviceId)|\(inboxesHex)|1"
        return Data(aadString.utf8)
    }
}