// 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.

package fr.tducray.mypwdtool.crypto

import java.security.SecureRandom
import java.util.Base64
import javax.crypto.AEADBadTagException
import javax.crypto.Cipher
import javax.crypto.SecretKeyFactory
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.PBEKeySpec
import javax.crypto.spec.SecretKeySpec

sealed class CryptoError(message: String) : Exception(message) {
    object KeyDerivationFailed : CryptoError("Key derivation failed")
    object EncryptionFailed : CryptoError("Encryption failed")
    object DecryptionFailed : CryptoError("Decryption failed — incorrect master password?")
    object InvalidBase64 : CryptoError("Invalid Base64 data")
    object InvalidData : CryptoError("Invalid data")
    object TamperedData : CryptoError("Corrupted or tampered data")
}

/**
 * JVM port of `Crypto/CryptoManager.swift` (the macOS/iOS app). Wire-format-compatible on
 * purpose, byte for byte — this is the single most important property of this file: an entry
 * encrypted on macOS/iOS must decrypt correctly here, and vice versa, since they share the same
 * vault via sync. See CryptoManagerCrossCompatTest for a test vector generated by the actual
 * Swift implementation, not just a Kotlin-only round-trip test (which would happily pass even if
 * this drifted from the Swift wire format in a way that broke real interop).
 *
 * Field-encryption blob (`encrypt`/`decrypt`): `[0x01][nonce:12][ciphertext][tag:16]`, base64.
 * DEK wrap (`wrapDEK`/`unwrapDEK`): `[nonce:12][ciphertext][tag:16]` — no leading version byte;
 * matches `AES.GCM.SealedBox.combined` used directly on the Swift side for that one, unlike the
 * field-encryption path which prepends `v1Version`.
 */
object CryptoManager {

    private const val V1_VERSION: Byte = 0x01
    private const val GCM_NONCE_LENGTH = 12
    private const val GCM_TAG_LENGTH_BITS = 128
    private const val GCM_TAG_LENGTH_BYTES = 16

    private val secureRandom = SecureRandom()

    // MARK: - Key derivation

    /** Derives a 256-bit KEK from the master password via PBKDF2-HMAC-SHA256. */
    fun deriveKey(password: String, salt: ByteArray, iterations: Int = 400_000): ByteArray {
        try {
            val factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256")
            val spec = PBEKeySpec(password.toCharArray(), salt, iterations, 256)
            return factory.generateSecret(spec).encoded
        } catch (e: Exception) {
            throw CryptoError.KeyDerivationFailed
        }
    }

    /** Generates a cryptographically random 32-byte salt. */
    fun generateSalt(): ByteArray = randomBytes(32)

    // MARK: - DEK

    fun generateDEK(): ByteArray = randomBytes(32)

    /** Wraps a DEK under a KEK using AES-GCM. Combined output: nonce(12) + ciphertext + tag(16). */
    fun wrapDEK(dek: ByteArray, kek: ByteArray): ByteArray = gcmSeal(plaintext = dek, key = kek, aad = null)

    /** Unwraps a DEK. Throws `DecryptionFailed` if the KEK is wrong or the data is tampered. */
    fun unwrapDEK(wrappedDEK: ByteArray, kek: ByteArray): ByteArray {
        try {
            return gcmOpen(combined = wrappedDEK, key = kek, aad = null)
        } catch (e: Exception) {
            throw CryptoError.DecryptionFailed
        }
    }

    // MARK: - Field encryption (v1)

    /** Encrypts plaintext with AAD. Blob: `[0x01][nonce:12][ciphertext][tag:16]`, base64-encoded. */
    fun encrypt(plaintext: String, key: ByteArray, aad: ByteArray): String {
        val data = plaintext.toByteArray(Charsets.UTF_8)
        val combined = gcmSeal(plaintext = data, key = key, aad = aad)
        val blob = ByteArray(1 + combined.size)
        blob[0] = V1_VERSION
        combined.copyInto(blob, destinationOffset = 1)
        return Base64.getEncoder().encodeToString(blob)
    }

    /** Decrypts a v1 blob produced by [encrypt]. */
    fun decrypt(base64Blob: String, key: ByteArray, aad: ByteArray): String {
        val blob = try {
            Base64.getDecoder().decode(base64Blob)
        } catch (e: IllegalArgumentException) {
            throw CryptoError.InvalidBase64
        }
        if (blob.size < 1 + GCM_NONCE_LENGTH + GCM_TAG_LENGTH_BYTES || blob[0] != V1_VERSION) {
            throw CryptoError.InvalidData
        }
        val combined = blob.copyOfRange(1, blob.size)
        try {
            val plainData = gcmOpen(combined = combined, key = key, aad = aad)
            return String(plainData, Charsets.UTF_8)
        } catch (e: CryptoError) {
            throw e
        } catch (e: Exception) {
            throw CryptoError.TamperedData
        }
    }

    // MARK: - AAD

    fun aad(vaultId: String, entryId: String, fieldName: String): ByteArray =
        "$vaultId|$entryId|$fieldName".toByteArray(Charsets.UTF_8)

    // MARK: - AES-GCM helpers

    private fun gcmSeal(plaintext: ByteArray, key: ByteArray, aad: ByteArray?): ByteArray {
        try {
            val nonce = randomBytes(GCM_NONCE_LENGTH)
            val cipher = Cipher.getInstance("AES/GCM/NoPadding")
            cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), GCMParameterSpec(GCM_TAG_LENGTH_BITS, nonce))
            if (aad != null) cipher.updateAAD(aad)
            val ciphertextAndTag = cipher.doFinal(plaintext)
            return nonce + ciphertextAndTag
        } catch (e: Exception) {
            throw CryptoError.EncryptionFailed
        }
    }

    /** Throws `AEADBadTagException` (via Cipher.doFinal) on a wrong key or tampered ciphertext. */
    private fun gcmOpen(combined: ByteArray, key: ByteArray, aad: ByteArray?): ByteArray {
        if (combined.size < GCM_NONCE_LENGTH + GCM_TAG_LENGTH_BYTES) throw CryptoError.InvalidData
        val nonce = combined.copyOfRange(0, GCM_NONCE_LENGTH)
        val ciphertextAndTag = combined.copyOfRange(GCM_NONCE_LENGTH, combined.size)
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(key, "AES"), GCMParameterSpec(GCM_TAG_LENGTH_BITS, nonce))
        if (aad != null) cipher.updateAAD(aad)
        return try {
            cipher.doFinal(ciphertextAndTag)
        } catch (e: AEADBadTagException) {
            throw CryptoError.TamperedData
        }
    }

    private fun randomBytes(count: Int): ByteArray {
        val bytes = ByteArray(count)
        secureRandom.nextBytes(bytes)
        return bytes
    }
}