package wang.yaojia.webterm.tlsandroid

import android.content.Context
import androidx.test.core.app.ApplicationProvider
import androidx.test.ext.junit.runners.AndroidJUnit4
import kotlinx.coroutines.runBlocking
import okhttp3.OkHttpClient
import org.junit.After
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertThrows
import org.junit.Assert.assertTrue
import org.junit.Before
import org.junit.Test
import org.junit.runner.RunWith
import wang.yaojia.webterm.clienttls.Pkcs12Parse

/**
 * Instrumented (real AndroidKeyStore + Tink — plan §7) tests for the re-reading mTLS seam.
 * COMPILES here; RUNS on a device/emulator during device QA.
 */
@RunWith(AndroidJUnit4::class)
class IdentityRepositoryTest {

    private val context: Context get() = ApplicationProvider.getApplicationContext()
    private val alias = "test-repo-identity"
    private val prefFile = "webterm_client_tls_test_prefs"
    private val keysetName = "webterm_cert_test_keyset"
    private val masterKeyUri = "android-keystore://webterm_cert_test_master_key"

    private lateinit var importer: AndroidKeyStoreImporter
    private lateinit var certStore: TinkCertStore
    private lateinit var client: OkHttpClient

    @Before
    fun setUp() {
        importer = AndroidKeyStoreImporter(alias = alias)
        certStore = TinkCertStore(context, keysetName, prefFile, masterKeyUri)
        client = OkHttpClient()
        wipe()
    }

    @After
    fun tearDown() = wipe()

    /** Both ping-pong slots + the pointer, so no state leaks across tests. */
    private fun wipe() {
        importer.remove(importer.primarySlot)
        importer.remove(importer.secondarySlot)
        certStore.clear()
    }

    private fun newRepository(store: CertStore = certStore): AndroidIdentityRepository =
        AndroidIdentityRepository(importer, store, client)

    @Test
    fun reReadingKeyManager_presentsMidRunImportedCertOnNextHandshake() {
        // A mutable holder stands in for the repository's live identity source.
        var installed: InstalledIdentity? = null
        val keyManager = ReReadingX509KeyManager { installed }

        // No identity → present nothing (a clean, classifiable handshake failure, never a wrong cert).
        assertNull(keyManager.chooseClientAlias(arrayOf("RSA"), null, null))
        assertNull(keyManager.getPrivateKey(alias))

        // Simulate a mid-run import (no relaunch, no factory rebuild).
        val parsed = Pkcs12Parse.parse(Fixtures.leafP12(), Fixtures.PASSPHRASE)
        installed = InstalledIdentity(
            alias = parsed.alias,
            keyAlgorithm = parsed.keyAlgorithm,
            privateKey = parsed.privateKey,
            leafCertificate = parsed.leafCertificate,
            issuerCertificates = parsed.issuerCertificates,
        )

        // The very next lookup re-reads the source and presents the freshly-imported identity.
        assertEquals(parsed.alias, keyManager.chooseClientAlias(arrayOf("RSA"), null, null))
        assertEquals(parsed.alias, keyManager.chooseEngineClientAlias(arrayOf("RSA"), null, null))
        assertNotNull(keyManager.getPrivateKey(parsed.alias))
        assertNotNull(keyManager.getCertificateChain(parsed.alias))
        // A foreign alias never leaks our key material.
        assertNull(keyManager.getPrivateKey("someone-else"))
    }

    @Test
    fun importPublishesIdentityAndSummary() = runBlocking {
        val repository = newRepository()
        assertFalse(repository.hasInstalledIdentity())
        assertNotNull("SSL material installed even with no cert", repository.sslMaterial().sslSocketFactory)

        val summary = repository.importIdentity(Fixtures.leafP12(), Fixtures.PASSPHRASE)

        assertEquals(Fixtures.LEAF_SUBJECT_CN, summary.subjectCommonName)
        assertEquals(Fixtures.LEAF_ISSUER_CN, summary.issuerCommonName)
        assertTrue(repository.hasInstalledIdentity())
        assertEquals(Fixtures.LEAF_SUBJECT_CN, repository.currentSummary()?.subjectCommonName)
        // First install lands on the primary slot; the pointer names it (single source of truth).
        assertEquals(importer.primarySlot, certStore.load()?.keyStoreAlias)
    }

    /**
     * FIX 3 (cache freshness): a device cert committed OUT OF BAND of a running repository (the zero-`.p12`
     * [DeviceEnroller] writes the leaf straight into the shared [CertStore] + AndroidKeyStore) is picked up
     * by [AndroidIdentityRepository.refreshFromStore] WITHOUT a process restart — the cached "no identity"
     * flips to the freshly-committed leaf and is presented on the next handshake.
     */
    @Test
    fun refreshFromStore_publishesAnOutOfBandCommittedIdentity_withoutRestart() = runBlocking {
        val running = newRepository()
        // Touch it while nothing is installed — caches the (null) initial identity.
        assertFalse(running.hasInstalledIdentity())

        // Simulate DeviceEnroller committing an identity out of band (a second repo over the SAME stores).
        newRepository().importIdentity(Fixtures.leafP12(), Fixtures.PASSPHRASE)

        // The running repo still shows its stale cache (no restart yet).
        assertFalse("stale cache still reports no identity before a refresh", running.hasInstalledIdentity())

        running.refreshFromStore()

        // The refresh re-read the committed live-pointer → the enrolled leaf is now live.
        assertTrue("refreshFromStore must publish the out-of-band committed identity", running.hasInstalledIdentity())
        assertEquals(Fixtures.LEAF_SUBJECT_CN, running.currentSummary()?.subjectCommonName)
    }

    @Test
    fun rotateThenRemove_evictsPooledConnections_andClearsIdentity() = runBlocking {
        val repository = newRepository()
        repository.importIdentity(Fixtures.leafP12(), Fixtures.PASSPHRASE)
        assertTrue(repository.hasInstalledIdentity())

        // Rotation stages into the OTHER slot and evicts pooled/resumed connections.
        repository.rotate(Fixtures.leafP12(), Fixtures.PASSPHRASE)
        assertTrue(repository.hasInstalledIdentity())
        assertEquals(importer.secondarySlot, certStore.load()?.keyStoreAlias)
        assertEquals(0, client.connectionPool.connectionCount())

        // Removal clears both stores and evicts again so nothing reuses the old identity.
        // (A populated pool needs a live TLS server → the no-reuse proof is the A34 E2E; here we
        //  assert eviction ran and left the pool empty + the identity cleared.)
        repository.remove()
        assertFalse(repository.hasInstalledIdentity())
        assertNull(repository.currentSummary())
        assertNull(certStore.load())
        assertEquals(0, client.connectionPool.connectionCount())
    }

    /**
     * The single-commit invariant (plan §8, A11): a rotation that fails AT the commit step must leave
     * the prior identity fully live and the two stores consistent (no half-applied new identity), and
     * a later successful install must fully replace it.
     */
    @Test
    fun failedCommit_keepsPriorIdentityLive_thenSuccessfulInstallReplacesIt() = runBlocking {
        val failing = FailingOnceCertStore(certStore)
        val repository = newRepository(failing)

        // Install identity A (lands on the primary slot; pointer names it).
        repository.importIdentity(Fixtures.leafP12(), Fixtures.PASSPHRASE)
        val slotA = certStore.load()!!.keyStoreAlias
        assertEquals(importer.primarySlot, slotA)
        assertNotNull("A's key present after install", importer.loadPrivateKey(slotA))

        // Attempt to rotate to B, injecting a failure at the COMMIT (metadata save) step.
        failing.failNextSave = true
        assertThrows(InjectedCommitFailure::class.java) {
            runBlocking { repository.rotate(Fixtures.leafP12(), Fixtures.PASSPHRASE) }
        }

        // (a) prior intact → the re-reading KeyManager (reads the same live view) still presents A.
        assertTrue("prior identity still live after a failed rotation", repository.hasInstalledIdentity())
        assertEquals(Fixtures.LEAF_SUBJECT_CN, repository.currentSummary()?.subjectCommonName)
        // (b) stores consistent: the pointer still names A's slot; the staged slot holds NO half-B key.
        assertEquals(slotA, certStore.load()!!.keyStoreAlias)
        assertNotNull("A's key untouched by the failed rotation", importer.loadPrivateKey(slotA))
        assertNull("no half-written B key in the staging slot", importer.loadPrivateKey(importer.secondarySlot))

        // (c) a subsequent successful install of B fully replaces A.
        failing.failNextSave = false
        repository.rotate(Fixtures.leafP12(), Fixtures.PASSPHRASE)
        assertTrue(repository.hasInstalledIdentity())
        assertEquals("pointer flipped to the other slot", importer.secondarySlot, certStore.load()!!.keyStoreAlias)
        assertNotNull("B's key now present", importer.loadPrivateKey(importer.secondarySlot))
        assertNull("A's superseded key GC'd", importer.loadPrivateKey(slotA))
    }

    /**
     * Regression for the `currentLive()` single-source-of-truth bug: after [remove], the repository
     * must report NO identity even when the lazily-loaded `initialIdentity` was already resolved to a
     * real value at startup (the normal relaunch-then-remove flow). The old `liveOverride?.value ?:
     * initialIdentity` collapsed the post-remove `Box(null)` back to the stale cached identity, so a
     * removed cert kept being reported live and would be presented in a handshake with a dangling key.
     */
    @Test
    fun remove_afterStartupTouch_reportsNoIdentity_notStaleCached() = runBlocking {
        // Persist identity A, then simulate an app relaunch: a FRESH repository over the same stores.
        newRepository().importIdentity(Fixtures.leafP12(), Fixtures.PASSPHRASE)
        val relaunched = newRepository()

        // Startup touch: resolve the lazy `initialIdentity` to the real installed identity A.
        assertTrue("identity restored from storage on relaunch", relaunched.hasInstalledIdentity())
        assertNotNull(relaunched.currentSummary())

        // Remove — must win over the already-cached initialIdentity.
        relaunched.remove()

        assertFalse("removed identity must NOT still report as installed", relaunched.hasInstalledIdentity())
        assertNull("removed identity must have no summary", relaunched.currentSummary())
        assertNull("pointer cleared", certStore.load())
        // The KeyManager reads the same live view → presents nothing after removal.
        assertNull(ReReadingX509KeyManager { null }.chooseClientAlias(arrayOf("RSA"), null, null))
    }

    /** Wraps a real [CertStore], failing the next [save] exactly once so the commit step can be forced to throw. */
    private class FailingOnceCertStore(private val delegate: CertStore) : CertStore {
        @Volatile
        var failNextSave: Boolean = false

        override fun save(metadata: StoredIdentityMetadata) {
            if (failNextSave) {
                failNextSave = false
                throw InjectedCommitFailure()
            }
            delegate.save(metadata)
        }

        override fun load(): StoredIdentityMetadata? = delegate.load()
        override fun clear(): Unit = delegate.clear()
    }

    private class InjectedCommitFailure : RuntimeException("injected commit failure")
}
