package app.closer.crypto

import com.google.crypto.tink.Aead
import com.google.crypto.tink.CleartextKeysetHandle
import com.google.crypto.tink.JsonKeysetReader
import com.google.crypto.tink.JsonKeysetWriter
import com.google.crypto.tink.KeysetHandle
import com.google.crypto.tink.aead.AeadKeyTemplates
import java.io.ByteArrayOutputStream
import java.util.Base64
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Encrypts and decrypts sealed answer payloads using a per-answer one-time AES-256-GCM key.
 *
 * The payload (writtenText, selectedOptionIds, scaleValue) is bundled as canonical JSON and
 * encrypted with Tink AEAD. The result is prefixed "sealed:v1:" to distinguish sealed docs
 * from legacy "enc:v1:" couple-key encrypted docs.
 *
 * AAD = "{coupleId}|{questionId}|{userId}" — binds the ciphertext to its origin so a
 * ciphertext cannot be transplanted to a different user or question.
 *
 * The one-time key itself never leaves the device before reveal. It is kept in
 * [PendingAnswerKeyStore] and only released (wrapped with the partner's public key) after
 * both partners have submitted.
 */
@Singleton
class SealedAnswerEncryptor @Inject constructor() {

    data class AnswerPayload(
        val writtenText: String?,
        val selectedOptionIds: List<String>,
        val scaleValue: Int?
    )

    fun generateOneTimeKey(): KeysetHandle =
        KeysetHandle.generateNew(AeadKeyTemplates.AES256_GCM)

    /**
     * Encrypts [payload] and returns the "sealed:v1:..." ciphertext string.
     */
    fun seal(
        payload: AnswerPayload,
        keyHandle: KeysetHandle,
        coupleId: String,
        questionId: String,
        userId: String
    ): String {
        val aead = keyHandle.getPrimitive(Aead::class.java)
        val plaintext = encodePayload(payload)
        val ciphertext = aead.encrypt(plaintext, aad(coupleId, questionId, userId))
        return SEALED_PREFIX + Base64.getUrlEncoder().withoutPadding().encodeToString(ciphertext)
    }

    /**
     * Decrypts a "sealed:v1:..." string back to [AnswerPayload].
     * Throws if the ciphertext is corrupt or the AAD does not match.
     */
    fun open(
        encryptedPayload: String,
        keyHandle: KeysetHandle,
        coupleId: String,
        questionId: String,
        userId: String
    ): AnswerPayload {
        require(encryptedPayload.startsWith(SEALED_PREFIX)) { "Not a sealed payload" }
        val ciphertext = Base64.getUrlDecoder().decode(encryptedPayload.removePrefix(SEALED_PREFIX))
        val aead = keyHandle.getPrimitive(Aead::class.java)
        val plaintext = aead.decrypt(ciphertext, aad(coupleId, questionId, userId))
        return decodePayload(plaintext)
    }

    fun serializeKey(keyHandle: KeysetHandle): String {
        val baos = ByteArrayOutputStream()
        CleartextKeysetHandle.write(keyHandle, JsonKeysetWriter.withOutputStream(baos))
        return baos.toString(Charsets.UTF_8.name())
    }

    fun deserializeKey(json: String): KeysetHandle =
        CleartextKeysetHandle.read(JsonKeysetReader.withString(json))

    private fun aad(coupleId: String, questionId: String, userId: String): ByteArray =
        "$coupleId|$questionId|$userId".toByteArray(Charsets.UTF_8)

    private fun encodePayload(payload: AnswerPayload): ByteArray {
        val ids = payload.selectedOptionIds.sorted().joinToString(",") { "\"${escape(it)}\"" }
        val text = if (payload.writtenText != null) "\"${escape(payload.writtenText)}\"" else "null"
        val scale = payload.scaleValue?.toString() ?: "null"
        val json = "{\"scaleValue\":$scale,\"selectedOptionIds\":[$ids],\"writtenText\":$text}"
        return json.toByteArray(Charsets.UTF_8)
    }

    private fun decodePayload(bytes: ByteArray): AnswerPayload {
        val json = bytes.toString(Charsets.UTF_8)
        // Manual parse — no Android JSON dependency so this runs in unit tests without Robolectric.
        val scaleValue = extractField(json, "scaleValue")?.toIntOrNull()
        val writtenText = extractString(json, "writtenText")
        val selectedOptionIds = extractArray(json, "selectedOptionIds")
        return AnswerPayload(writtenText, selectedOptionIds, scaleValue)
    }

    private fun extractField(json: String, key: String): String? {
        val pattern = "\"$key\":([^,}]+)".toRegex()
        return pattern.find(json)?.groupValues?.get(1)?.trim()?.takeIf { it != "null" }
    }

    private fun extractString(json: String, key: String): String? {
        val pattern = "\"$key\":\"((?:[^\"\\\\]|\\\\.)*)\"".toRegex()
        return pattern.find(json)?.groupValues?.get(1)?.unescape()
    }

    private fun extractArray(json: String, key: String): List<String> {
        val pattern = "\"$key\":\\[([^]]*)]".toRegex()
        val inner = pattern.find(json)?.groupValues?.get(1) ?: return emptyList()
        if (inner.isBlank()) return emptyList()
        return "\"((?:[^\"\\\\]|\\\\.)*)\"".toRegex()
            .findAll(inner)
            .map { it.groupValues[1].unescape() }
            .toList()
    }

    private fun escape(s: String): String = buildString {
        for (c in s) when (c) {
            '\\' -> append("\\\\")
            '"'  -> append("\\\"")
            '\n' -> append("\\n")
            '\r' -> append("\\r")
            '\t' -> append("\\t")
            else -> append(c)
        }
    }

    private fun String.unescape(): String = replace("\\\"", "\"")
        .replace("\\\\", "\\")
        .replace("\\n", "\n")
        .replace("\\r", "\r")
        .replace("\\t", "\t")

    companion object {
        const val SEALED_PREFIX = "sealed:v1:"
    }
}
