package ch.threema.domain.protocol.multidevice

import ch.threema.base.ThreemaException
import ch.threema.base.crypto.NaCl
import ch.threema.base.crypto.Nonce
import ch.threema.base.utils.getThreemaLogger
import ch.threema.common.emptyByteArray
import ch.threema.common.generateRandomBytes
import ch.threema.common.secureRandom
import ch.threema.domain.protocol.connection.data.D2dMessage
import ch.threema.domain.protocol.connection.data.D2mProtocolException
import ch.threema.domain.protocol.connection.data.InboundD2mMessage
import ch.threema.libthreema.CryptoException
import ch.threema.libthreema.blake2bMac256
import ch.threema.protobuf.d2d.DeviceInfo
import ch.threema.protobuf.d2d.Envelope
import ch.threema.protobuf.d2d.TransactionScope
import ch.threema.protobuf.d2d.transactionScope

private val logger = getThreemaLogger("MultiDeviceKeys")

data class MultiDeviceKeys(val dgk: ByteArray) {
    companion object {
        private const val BLAKE_2B_PERSONAL = "3ma-mdev"
        private const val SALT_DGPK = "p"
        private const val SALT_DGRK = "r"
        private const val SALT_DGDIK = "di"
        private const val SALT_DGSDDK = "sdd"
        private const val SALT_DGTSK = "ts"

        private fun createNonce(): ByteArray =
            secureRandom().generateRandomBytes(NaCl.NONCE_BYTES)

        @Throws(ThreemaException::class)
        private fun blake2b256(
            key: ByteArray,
            salt: String,
        ): ByteArray =
            try {
                blake2bMac256(
                    key = key,
                    personal = BLAKE_2B_PERSONAL.encodeToByteArray(),
                    salt = salt.encodeToByteArray(),
                    data = emptyByteArray(),
                )
            } catch (cryptoException: CryptoException.InvalidParameter) {
                logger.error("Failed to compute blake2b hash", cryptoException)
                throw Error("Failed to compute blake2b hash", cryptoException)
            }
    }

    init {
        // assert correct length of dgk
        if (dgk.size != 32) {
            throw D2mProtocolException("Invalid length of key material. Expected 32, actual ${dgk.size}")
        }
    }

    internal val dgpk: ByteArray by lazy { blake2b256(key = dgk, salt = SALT_DGPK) }
    val dgid: ByteArray by lazy { NaCl.derivePublicKey(dgpk) }

    internal val dgrk: ByteArray by lazy { blake2b256(key = dgk, salt = SALT_DGRK) }
    internal val dgdik: ByteArray by lazy { blake2b256(key = dgk, salt = SALT_DGDIK) }
    internal val dgsddk: ByteArray by lazy { blake2b256(key = dgk, salt = SALT_DGSDDK) }
    internal val dgtsk: ByteArray by lazy { blake2b256(key = dgk, salt = SALT_DGTSK) }

    internal fun createServerHelloResponse(serverHello: InboundD2mMessage.ServerHello): ByteArray {
        val naCl = NaCl(dgpk, serverHello.esk)
        val nonce = createNonce()
        return nonce + naCl.encrypt(data = serverHello.challenge, nonce = nonce)
    }

    fun encryptDeviceInfo(deviceInfo: D2dMessage.DeviceInfo): ByteArray {
        val nonce = createNonce()
        return nonce + NaCl.symmetricEncryptData(data = deviceInfo.bytes, key = dgdik, nonce = nonce)
    }

    fun decryptDeviceInfo(deviceInfo: ByteArray): D2dMessage.DeviceInfo {
        val nonce = deviceInfo.copyOfRange(0, NaCl.NONCE_BYTES)
        val data = deviceInfo.copyOfRange(nonce.size, deviceInfo.size)
        val decrypted = NaCl.symmetricDecryptData(
            data = data,
            key = dgdik,
            nonce = nonce,
        )
        return DeviceInfo.parseFrom(decrypted)
            .let { D2dMessage.DeviceInfo.fromProtobuf(it) }
    }

    fun encryptTransactionScope(scope: TransactionScope.Scope): ByteArray {
        val nonce = createNonce()
        val bytes = transactionScope {
            this.scope = scope
        }.toByteArray()
        val encrypted = NaCl.symmetricEncryptData(data = bytes, key = dgtsk, nonce = nonce)
        return nonce + encrypted
    }

    fun decryptTransactionScope(encryptedTransactionScope: ByteArray): TransactionScope.Scope {
        val nonce = encryptedTransactionScope.copyOfRange(0, NaCl.NONCE_BYTES)
        val data = encryptedTransactionScope.copyOfRange(nonce.size, encryptedTransactionScope.size)
        val decrypted = NaCl.symmetricDecryptData(
            data = data,
            key = dgtsk,
            nonce = nonce,
        )
        return TransactionScope.parseFrom(decrypted).scope
    }

    fun encryptEnvelope(envelope: Envelope): EncryptedEnvelopeResult {
        val nonceBytes = createNonce()
        val encryptedEnvelope = nonceBytes + NaCl.symmetricEncryptData(data = envelope.toByteArray(), key = dgrk, nonce = nonceBytes)
        return EncryptedEnvelopeResult(
            encryptedEnvelope = encryptedEnvelope,
            nonce = Nonce(nonceBytes),
            debugInfo = EncryptedEnvelopeResult.DebugInfo(
                protoContentCaseName = envelope.contentCase.name,
                rawEnvelopeContent = envelope.toString(),
            ),
        )
    }

    fun decryptEnvelope(encryptedEnvelopeBytes: ByteArray): Pair<Nonce, Envelope> {
        val nonce = encryptedEnvelopeBytes.copyOfRange(0, NaCl.NONCE_BYTES)
        val data = encryptedEnvelopeBytes.copyOfRange(nonce.size, encryptedEnvelopeBytes.size)
        val decrypted = NaCl.symmetricDecryptData(
            data = data,
            key = dgrk,
            nonce = nonce,
        )
        return Nonce(nonce) to Envelope.parseFrom(decrypted)
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (other !is MultiDeviceKeys) return false

        if (!dgk.contentEquals(other.dgk)) return false

        return true
    }

    override fun hashCode(): Int {
        return dgk.contentHashCode()
    }

    override fun toString(): String {
        return "MultiDeviceKeys(dgk=********)"
    }

    /**
     * @param encryptedEnvelope Encrypted envelope bytes
     * @param debugInfo         Only used for debugging. Contains the unencrypted contents of [Envelope.toString]
     */
    data class EncryptedEnvelopeResult(
        val encryptedEnvelope: ByteArray,
        val nonce: Nonce,
        val debugInfo: DebugInfo,
    ) {
        override fun equals(other: Any?): Boolean {
            if (this === other) return true
            if (javaClass != other?.javaClass) return false
            other as EncryptedEnvelopeResult
            if (!encryptedEnvelope.contentEquals(other.encryptedEnvelope)) return false
            if (nonce != other.nonce) return false
            if (debugInfo != other.debugInfo) return false
            return true
        }

        override fun hashCode(): Int {
            var result = encryptedEnvelope.contentHashCode()
            result = 31 * result + nonce.hashCode()
            result = 31 * result + debugInfo.hashCode()
            return result
        }

        override fun toString(): String {
            return "EncryptedEnvelopeResult(encryptedEnvelope: ${encryptedEnvelope.contentToString()}, nonce: ***, debugInfo: $debugInfo"
        }

        /**
         * @param protoContentCaseName Is the value of [Envelope.getContentCase]
         * @param rawEnvelopeContent   Contains the whole proto message contents in an **unencrypted** form. Never send this
         * and only log it on debug level!
         */
        data class DebugInfo(
            val protoContentCaseName: String,
            val rawEnvelopeContent: String,
        ) {
            override fun toString(): String {
                return "EncryptedEnvelopeResult(protoContentCaseName: $protoContentCaseName, rawEnvelopeContent: ***)"
            }
        }
    }
}
