package ch.threema.localcrypto

import ch.threema.base.utils.toByteString
import ch.threema.common.toByteArray
import ch.threema.localcrypto.models.Argon2Version
import ch.threema.localcrypto.models.MasterKeyStorageData
import ch.threema.localcrypto.models.Version2MasterKeyStorageInnerData
import ch.threema.localcrypto.models.Version2MasterKeyStorageOuterData
import ch.threema.localcrypto.protobuf.InnerKeyStorage
import ch.threema.localcrypto.protobuf.IntermediateKeyStorage
import ch.threema.localcrypto.protobuf.IntermediateKeyStorageV1Kt.remoteSecretProtected
import ch.threema.localcrypto.protobuf.OuterKeyStorage
import ch.threema.localcrypto.protobuf.OuterKeyStorageV1.Argon2idProtected
import ch.threema.localcrypto.protobuf.OuterKeyStorageV1Kt
import ch.threema.localcrypto.protobuf.innerKeyStorageV1
import ch.threema.localcrypto.protobuf.intermediateKeyStorageV1
import ch.threema.localcrypto.protobuf.outerKeyStorageV1
import ch.threema.protobuf.combineEncryptedDataAndNonce
import com.google.protobuf.ByteString
import com.google.protobuf.Internal
import com.google.protobuf.MessageLite
import com.google.protobuf.kotlin.plus
import com.google.protobuf.kotlin.toByteString
import java.nio.ByteOrder

class Version2MasterKeyStorageEncoder {
    fun encodeMasterKeyStorageData(masterKeyStorageData: MasterKeyStorageData.Version2): ByteString =
        encodeOuterData(masterKeyStorageData.outerData)

    private fun encodeOuterData(outerData: Version2MasterKeyStorageOuterData): ByteString =
        encodeVersionedMessage(
            version = OuterKeyStorage.Version.V1_0,
            message = outerKeyStorageV1 {
                when (outerData) {
                    is Version2MasterKeyStorageOuterData.PassphraseProtected -> {
                        argon2IdProtectedIntermediate = encodePassphraseProtected(outerData)
                    }

                    is Version2MasterKeyStorageOuterData.NotPassphraseProtected -> {
                        plaintextIntermediate = encodeInnerData(outerData.innerData)
                    }
                }
            },
        )

    private fun encodePassphraseProtected(outerData: Version2MasterKeyStorageOuterData.PassphraseProtected): Argon2idProtected =
        OuterKeyStorageV1Kt.argon2idProtected {
            version = when (outerData.argonVersion) {
                Argon2Version.VERSION_1_3 -> Argon2idProtected.Argon2Version.VERSION_1_3
            }
            iterations = outerData.iterations
            memoryBytes = outerData.memoryBytes
            parallelism = outerData.parallelism
            salt = outerData.salt.toByteString()
            encryptedIntermediate = combineEncryptedDataAndNonce(
                data = outerData.encryptedData.value,
                nonce = outerData.nonce.value,
            ).toByteString()
        }

    fun encodeInnerData(innerData: Version2MasterKeyStorageInnerData): ByteString =
        encodeVersionedMessage(
            version = IntermediateKeyStorage.Version.V1_0,
            message = intermediateKeyStorageV1 {
                when (innerData) {
                    is Version2MasterKeyStorageInnerData.RemoteSecretProtected -> {
                        this.remoteSecretProtectedInner = remoteSecretProtected {
                            remoteSecretHash = innerData.parameters.remoteSecretHash.toByteString()
                            remoteSecretAuthenticationToken = innerData.parameters.authenticationToken.toByteString()
                            encryptedInner = innerData.encryptedData.toByteString()
                        }
                    }
                    is Version2MasterKeyStorageInnerData.Unprotected -> {
                        plaintextInner = encodeVersionedMessage(
                            version = InnerKeyStorage.Version.V1_0,
                            message = innerKeyStorageV1 {
                                masterKey = innerData.masterKeyData.toByteString()
                            },
                        )
                    }
                }
            },
        )

    private fun encodeVersionedMessage(version: Internal.EnumLite, message: MessageLite): ByteString =
        version.toByteString() + message.toByteString()

    private fun Internal.EnumLite.toByteString(): ByteString =
        number.toShort().toByteArray(order = ByteOrder.LITTLE_ENDIAN).toByteString()
}
