package wang.yaojia.webterm.tlsandroid

import android.content.Context
import android.content.SharedPreferences
import android.util.Base64
import com.google.crypto.tink.Aead
import com.google.crypto.tink.KeyTemplates
import com.google.crypto.tink.RegistryConfiguration
import com.google.crypto.tink.aead.AeadConfig
import com.google.crypto.tink.integration.android.AndroidKeysetManager
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.DataInputStream
import java.io.DataOutputStream

/**
 * B4 · The auxiliary enrollment record needed to drive silent rotation: the server-minted
 * [deviceId] (the `/device/:id/renew` path segment), the [deviceName] re-used as the renew CSR
 * subject CN, the AndroidKeyStore [keyStoreAlias] holding the SAME non-exportable key to re-sign
 * with, and [renewAfterEpochSeconds] (0 = unknown) for the rotation scheduler.
 *
 * This is deliberately SEPARATE from [StoredIdentityMetadata] (the mTLS live-pointer): the cert
 * identity is what the handshake presents; this record only exists so renew can find the device and
 * its key. The private key is never here — it stays non-exportable in AndroidKeyStore.
 */
public data class EnrollmentRecord(
    val deviceId: String,
    val deviceName: String,
    val keyStoreAlias: String,
    val renewAfterEpochSeconds: Long,
) {
    init {
        require(deviceId.isNotBlank()) { "deviceId must not be blank" }
        require(keyStoreAlias.isNotBlank()) { "keyStoreAlias must not be blank" }
    }
}

/** Length-prefixed binary codec for [EnrollmentRecord] (KISS — three UTF strings + one long). */
public object EnrollmentRecordCodec {
    public fun encode(record: EnrollmentRecord): ByteArray {
        val out = ByteArrayOutputStream()
        DataOutputStream(out).use { data ->
            data.writeUTF(record.deviceId)
            data.writeUTF(record.deviceName)
            data.writeUTF(record.keyStoreAlias)
            data.writeLong(record.renewAfterEpochSeconds)
        }
        return out.toByteArray()
    }

    /** Decode [bytes]; any structural failure → [CorruptStoredIdentityException]. */
    public fun decode(bytes: ByteArray): EnrollmentRecord =
        try {
            DataInputStream(ByteArrayInputStream(bytes)).use { data ->
                EnrollmentRecord(
                    deviceId = data.readUTF(),
                    deviceName = data.readUTF(),
                    keyStoreAlias = data.readUTF(),
                    renewAfterEpochSeconds = data.readLong(),
                )
            }
        } catch (e: Exception) {
            throw CorruptStoredIdentityException("Stored enrollment record was truncated/malformed", e)
        }
}

/**
 * Storage contract for the [EnrollmentRecord] — repository pattern so [DeviceEnroller] depends on
 * the operation set and a fault/blank can be injected in tests. Idempotent [clear].
 */
public interface EnrollmentRecordStore {
    public fun save(record: EnrollmentRecord)

    public fun load(): EnrollmentRecord?

    public fun clear()
}

/**
 * Tink-AEAD-encrypted [EnrollmentRecordStore] over an app-private `SharedPreferences` file, mirroring
 * [TinkCertStore]'s custody model (AndroidKeystore-wrapped master key; uninstall-wiped; useless off
 * this device). Kept in its own key/file namespace so it never collides with the cert live-pointer.
 */
public class TinkEnrollmentRecordStore(
    context: Context,
    private val keysetName: String = DEFAULT_KEYSET_NAME,
    private val prefFileName: String = DEFAULT_PREF_FILE,
    private val masterKeyUri: String = DEFAULT_MASTER_KEY_URI,
) : EnrollmentRecordStore {
    private val appContext: Context = context.applicationContext
    private val aead: Aead by lazy { buildAead() }

    override fun save(record: EnrollmentRecord) {
        val ciphertext = aead.encrypt(EnrollmentRecordCodec.encode(record), ASSOCIATED_DATA)
        val committed = prefs().edit()
            .putString(BLOB_KEY, Base64.encodeToString(ciphertext, Base64.NO_WRAP))
            .commit()
        if (!committed) throw java.io.IOException("Failed to durably persist the device enrollment record")
    }

    override fun load(): EnrollmentRecord? {
        val encoded = prefs().getString(BLOB_KEY, null) ?: return null
        val ciphertext = try {
            Base64.decode(encoded, Base64.NO_WRAP)
        } catch (e: IllegalArgumentException) {
            throw CorruptStoredIdentityException("Enrollment record blob was not valid base64", e)
        }
        val plaintext = try {
            aead.decrypt(ciphertext, ASSOCIATED_DATA)
        } catch (e: java.security.GeneralSecurityException) {
            throw CorruptStoredIdentityException("Enrollment record blob failed AEAD decryption", e)
        }
        return EnrollmentRecordCodec.decode(plaintext)
    }

    override fun clear() {
        val committed = prefs().edit().remove(BLOB_KEY).commit()
        if (!committed) throw java.io.IOException("Failed to durably clear the device enrollment record")
    }

    private fun buildAead(): Aead {
        AeadConfig.register()
        val keysetHandle = AndroidKeysetManager.Builder()
            .withSharedPref(appContext, keysetName, prefFileName)
            .withKeyTemplate(KeyTemplates.get(AEAD_KEY_TEMPLATE))
            .withMasterKeyUri(masterKeyUri)
            .build()
            .keysetHandle
        return keysetHandle.getPrimitive(RegistryConfiguration.get(), Aead::class.java)
    }

    private fun prefs(): SharedPreferences =
        appContext.getSharedPreferences(prefFileName, Context.MODE_PRIVATE)

    public companion object {
        private const val DEFAULT_KEYSET_NAME = "webterm_enroll_keyset"
        private const val DEFAULT_PREF_FILE = "webterm_enroll_record_prefs"
        private const val DEFAULT_MASTER_KEY_URI = "android-keystore://webterm_enroll_master_key"
        private const val AEAD_KEY_TEMPLATE = "AES256_GCM"
        private const val BLOB_KEY = "enrollment_record_blob"
        private val ASSOCIATED_DATA: ByteArray = "webterm.client-tls.enrollment-record".toByteArray()
    }
}
