package com.tomhempel.rbarometer.security

import android.util.Base64
import com.google.gson.JsonObject
import com.google.gson.JsonParser
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.Mac
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * Session-token-derived envelope encryption for sensitive mobile API payloads.
 */
object SessionPayloadCrypto {

    private const val HKDF_INFO = "qdb-aes"
    private const val KEY_LEN = 32
    private const val ALG_GCM = "A256GCM"
    private const val GCM_TAG_BITS = 128

    fun envelope(plainJson: String, tokenHex: String): JsonObject {
        val encrypted = encryptGcm(plainJson.toByteArray(Charsets.UTF_8), sessionKey(tokenHex))
        return JsonObject().apply {
            addProperty("encrypted", true)
            addProperty("alg", ALG_GCM)
            addProperty("payload", Base64.encodeToString(encrypted, Base64.NO_WRAP))
        }
    }

    fun decryptEnvelope(envelope: JsonObject, tokenHex: String): JsonObject {
        if (!envelope.get("encrypted")?.asBoolean.orFalse()) {
            throw IllegalArgumentException("Expected encrypted envelope")
        }
        val payloadB64 = envelope.get("payload")?.asString
            ?: throw IllegalArgumentException("Missing payload")
        val packed = Base64.decode(payloadB64, Base64.NO_WRAP)
        val plain = when (envelope.get("alg")?.asString) {
            ALG_GCM -> decryptGcm(packed, sessionKey(tokenHex))
            else -> decryptCbc(packed, sessionKey(tokenHex))
        }
        return JsonParser.parseString(String(plain, Charsets.UTF_8)).asJsonObject
    }

    fun decryptDataElement(data: JsonObject?, tokenHex: String): JsonObject {
        if (data == null) return JsonObject()
        return if (data.get("encrypted")?.asBoolean == true) {
            decryptEnvelope(data, tokenHex)
        } else {
            data
        }
    }

    private fun sessionKey(tokenHex: String): ByteArray {
        val ikm = hexToBytes(tokenHex.trim())
        return hkdfSha256(ikm, HKDF_INFO.toByteArray(Charsets.UTF_8), KEY_LEN)
    }

    private fun encryptGcm(plain: ByteArray, key: ByteArray): ByteArray {
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        val iv = ByteArray(12).also { SecureRandom().nextBytes(it) }
        cipher.init(
            Cipher.ENCRYPT_MODE,
            SecretKeySpec(normalizeKey(key), "AES"),
            GCMParameterSpec(GCM_TAG_BITS, iv),
        )
        val ct = cipher.doFinal(plain)
        return iv + ct
    }

    private fun decryptGcm(data: ByteArray, key: ByteArray): ByteArray {
        val ivLen = 12
        if (data.size <= ivLen) throw IllegalArgumentException("Ciphertext too short")
        val iv = data.copyOfRange(0, ivLen)
        val ct = data.copyOfRange(ivLen, data.size)
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(
            Cipher.DECRYPT_MODE,
            SecretKeySpec(normalizeKey(key), "AES"),
            GCMParameterSpec(GCM_TAG_BITS, iv),
        )
        return cipher.doFinal(ct)
    }

    private fun decryptCbc(data: ByteArray, key: ByteArray): ByteArray {
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
        if (data.size < 16) throw IllegalArgumentException("Ciphertext too short")
        val iv = data.copyOfRange(0, 16)
        val ct = data.copyOfRange(16, data.size)
        cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(normalizeKey(key), "AES"), IvParameterSpec(iv))
        return cipher.doFinal(ct)
    }

    private fun normalizeKey(key: ByteArray): ByteArray {
        val out = ByteArray(KEY_LEN)
        val copyLen = minOf(key.size, KEY_LEN)
        key.copyInto(out, 0, 0, copyLen)
        return out
    }

    private fun hkdfSha256(ikm: ByteArray, info: ByteArray, length: Int): ByteArray {
        val salt = ByteArray(32)
        val prkMac = Mac.getInstance("HmacSHA256")
        prkMac.init(SecretKeySpec(salt, "HmacSHA256"))
        val prk = prkMac.doFinal(ikm)
        val okm = ByteArray(length)
        var t = ByteArray(0)
        var offset = 0
        var counter = 1
        while (offset < length) {
            val mac = Mac.getInstance("HmacSHA256")
            mac.init(SecretKeySpec(prk, "HmacSHA256"))
            mac.update(t)
            mac.update(info)
            mac.update(counter.toByte())
            t = mac.doFinal()
            val copy = minOf(t.size, length - offset)
            t.copyInto(okm, offset, 0, copy)
            offset += copy
            counter++
        }
        return okm
    }

    private fun hexToBytes(hex: String): ByteArray {
        val clean = hex.trim()
        if (clean.length % 2 != 0) return clean.toByteArray(Charsets.UTF_8)
        return ByteArray(clean.length / 2) { i ->
            clean.substring(i * 2, i * 2 + 2).toInt(16).toByte()
        }
    }

    private fun Boolean?.orFalse() = this == true
}
