package wang.yaojia.webterm.tlsandroid

import android.util.Log
import java.security.KeyStore
import javax.net.ssl.KeyManager
import javax.net.ssl.SSLContext
import javax.net.ssl.SSLSocketFactory
import javax.net.ssl.TrustManager
import javax.net.ssl.TrustManagerFactory
import javax.net.ssl.X509TrustManager
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import okhttp3.OkHttpClient
import wang.yaojia.webterm.clienttls.CertificateSummary
import wang.yaojia.webterm.clienttls.CertificateSummaryReader

/**
 * The `(SSLSocketFactory, X509TrustManager)` pair the shared `OkHttpClient` is built with. The
 * factory wraps the re-reading KeyManager (so it is STABLE across cert rotations — rotation is
 * handled inside the KeyManager, not by swapping the factory), and the trust manager is the default
 * system trust (tunnel/Tailscale present real LE certs — plan §6.9). This is the clean seam the `:app`
 * wiring (A15) bridges to `:transport-okhttp`'s `ClientIdentityProvider`; A11 never depends on that
 * module.
 */
public class ClientSslMaterial(
    public val sslSocketFactory: SSLSocketFactory,
    public val trustManager: X509TrustManager,
)

/**
 * The current device client identity, exposed as a re-reading mTLS seam (plan §2 mTLS row, §8, R4).
 *
 * The [sslMaterial] pair is built ONCE and installed on the shared client for its whole lifetime; a
 * mid-run [importIdentity]/[rotate] is presented on the NEXT handshake (re-reading KeyManager) and
 * [rotate]/[remove] additionally evict pooled/resumed connections so they can't reuse the old
 * identity. No relaunch, no client rebuild.
 */
public interface IdentityRepository {
    /** The stable SSL material to install on the shared `OkHttpClient` (always present — presents no
     *  client cert when no identity is installed). */
    public fun sslMaterial(): ClientSslMaterial

    /** Summary of the currently installed identity for the cert screen, or null if none installed. */
    public fun currentSummary(): CertificateSummary?

    /** Cheap gating check — is a device identity installed? */
    public fun hasInstalledIdentity(): Boolean

    /** Import a `.p12` (validates before persisting), then evict pooled connections. Returns the summary. */
    public suspend fun importIdentity(p12Bytes: ByteArray, passphrase: String): CertificateSummary

    /** Replace the installed identity with a new `.p12` (rotation — single-commit, prior stays live on failure). */
    public suspend fun rotate(p12Bytes: ByteArray, passphrase: String): CertificateSummary

    /** Remove the installed identity, then evict pooled connections so none reuse it. Idempotent. */
    public suspend fun remove()
}

/**
 * B4 · A narrow seam the zero-`.p12` enroll/renew commit ([DeviceEnroller]) fires so an in-memory
 * identity cache re-reads the freshly-committed live-pointer and presents the new leaf on the NEXT
 * mTLS handshake WITHOUT a process restart. Kept separate from [IdentityRepository] so the enroller
 * depends only on this one operation (it never needs the import/rotate/remove surface). The production
 * implementation is [AndroidIdentityRepository]; a JVM test uses a recording double.
 */
public fun interface IdentityCacheRefresher {
    /** Reload the persisted live identity into the in-memory cache and drop stale pooled connections. */
    public fun refreshFromStore()
}

/**
 * Default [IdentityRepository] over [AndroidKeyStoreImporter] (key home) + [CertStore] (encrypted
 * live-pointer at rest) + the shared [OkHttpClient] (for `connectionPool.evictAll()`).
 *
 * ### Single-commit rotation (the security invariant, plan §8, A11)
 * An import/rotation that fails at ANY point leaves the PRIOR working identity fully live and NEVER
 * leaves the two stores diverged. This is achieved by staging, then committing with one atomic write:
 *  1. Import the new key into the ping-pong slot that is NOT currently live
 *     ([AndroidKeyStoreImporter] `primary`/`secondary`) — the live key is never overwritten.
 *  2. COMMIT by persisting one blob to the [CertStore] whose [StoredIdentityMetadata.keyStoreAlias]
 *     names the staged slot. That single write is the source-of-truth "live pointer" flip.
 *  - Any failure BEFORE the flip → the prior pointer + key are untouched → prior identity fully live
 *    (the half-written staging key is best-effort deleted and is never read as live).
 *  - Any failure AFTER the flip → the NEW identity is fully live (the superseded key is best-effort
 *    GC'd; a stray unreferenced key is harmless and gets overwritten on the next rotation).
 * Startup and the re-reading KeyManager resolve the live identity SOLELY via that pointer.
 * Validate-before-persist still holds (a bad passphrase throws before any keystore/store mutation).
 *
 * ### Concurrency & lazy init
 * The mutating methods are `suspend` and serialized by a [Mutex] so overlapping calls (double-tap)
 * can never interleave the two-store commit. The initial on-disk identity is loaded LAZILY on first
 * access (not in the constructor), so this repository can be constructed on any thread with no I/O.
 *
 * A15 wiring note: the FIRST touch of the identity (e.g. [currentSummary]/[hasInstalledIdentity], a
 * mutator, or the first TLS handshake) performs synchronous AndroidKeyStore + Tink I/O — trigger that
 * first access off `Dispatchers.Main` (a background warm-up), never on the UI thread.
 *
 * The cached runtime identity always holds the AndroidKeyStore-backed (non-exportable) private-key
 * handle, never the transient parsed software key.
 */
public class AndroidIdentityRepository(
    private val importer: AndroidKeyStoreImporter,
    private val certStore: CertStore,
    private val sharedClient: OkHttpClient,
) : IdentityRepository, IdentityCacheRefresher {

    /** The live identity as tracked by the repo: the KeyManager view + which physical slot holds the key. */
    private class LiveIdentity(val installed: InstalledIdentity, val keyStoreAlias: String)

    /** Distinguishes "not yet mutated" (null) from "explicitly set, possibly to null" ([Box]). */
    private class Box<T>(val value: T)

    // Serializes the two-store commit so a double-tap can't interleave two rotations (FIX 2).
    private val commitMutex = Mutex()

    // Lazily loaded on FIRST access (never in the constructor) — no slow I/O to construct (FIX 3).
    private val initialIdentity: LiveIdentity? by lazy(LazyThreadSafetyMode.SYNCHRONIZED) {
        loadInstalledOrNull()
    }

    // Post-mutation state shadows the lazily-loaded initial value; volatile read on the handshake path.
    @Volatile
    private var liveOverride: Box<LiveIdentity?>? = null

    // Resolve via the Box PRESENCE, not the unwrapped value's null-ness: after remove() the override
    // is Box(null) (explicitly "no identity"), which MUST win over a still-cached non-null
    // initialIdentity — otherwise a removed cert (whose key/pointer are already deleted) would keep
    // being reported live and presented in a handshake with a dangling key handle. `?: initialIdentity`
    // would collapse Box(null) back to the stale identity, so this cannot use the Elvis operator.
    private fun currentLive(): LiveIdentity? {
        val override = liveOverride
        return if (override != null) override.value else initialIdentity
    }

    // Default system trust: tunnel/Tailscale present real LE certs (plan §6.9). Built once.
    private val systemTrustManager: X509TrustManager = systemDefaultTrustManager()

    // Stable factory wrapping the re-reading KeyManager over `currentLive()`. Built once.
    private val socketFactory: SSLSocketFactory = buildSocketFactory()

    override fun sslMaterial(): ClientSslMaterial =
        ClientSslMaterial(sslSocketFactory = socketFactory, trustManager = systemTrustManager)

    override fun currentSummary(): CertificateSummary? =
        currentLive()?.let { CertificateSummaryReader.summarize(it.installed.leafCertificate) }

    override fun hasInstalledIdentity(): Boolean = currentLive() != null

    override suspend fun importIdentity(p12Bytes: ByteArray, passphrase: String): CertificateSummary =
        install(p12Bytes, passphrase)

    override suspend fun rotate(p12Bytes: ByteArray, passphrase: String): CertificateSummary =
        install(p12Bytes, passphrase)

    override suspend fun remove(): Unit = commitMutex.withLock {
        // Clear the pointer FIRST (no live identity even if a key delete lags), then both key slots.
        certStore.clear()
        importer.remove(importer.primarySlot)
        importer.remove(importer.secondarySlot)
        liveOverride = Box(null)
        // Drop pooled/resumed connections so nothing reuses the removed identity (R4/§8).
        sharedClient.connectionPool.evictAll()
    }

    /**
     * FIX 3 (cache freshness) · Re-read the persisted live-pointer into the in-memory cache. Used when a
     * device certificate is committed OUT OF BAND of this repository — the zero-`.p12` [DeviceEnroller]
     * writes the leaf straight into the shared [CertStore] + AndroidKeyStore, so without this the running
     * repo would keep presenting its cached (pre-enroll) identity until process restart. Publishing the
     * freshly-loaded snapshot as [liveOverride] and evicting pooled connections makes the enrolled leaf
     * present on the NEXT handshake. Reloading to `null` (a fault/absent pointer) is a valid outcome and
     * simply reports "no identity". Not `suspend` — the enroller already runs this off the UI thread.
     */
    override fun refreshFromStore() {
        liveOverride = Box(loadInstalledOrNull())
        sharedClient.connectionPool.evictAll()
    }

    /**
     * Single-commit install/rotation (see the class KDoc). Validation throws before any mutation;
     * the new key is imported into the non-live slot; the COMMIT is one atomic [CertStore.save] that
     * flips the live pointer. Only after the flip is the new identity published and the superseded
     * key GC'd. Any earlier failure leaves the prior identity fully live and the stores consistent.
     */
    private suspend fun install(p12Bytes: ByteArray, passphrase: String): CertificateSummary =
        commitMutex.withLock {
            val liveAlias = currentLive()?.keyStoreAlias
            val stagingAlias = stagingSlotFor(liveAlias)

            // Stage into the NON-live slot then COMMIT. Everything up to and including the atomic
            // certStore.save() flip is wrapped so ANY pre-flip failure (import, key readback, encode,
            // or the commit write) best-effort drops the half-written staging key and rethrows —
            // leaving the prior pointer + live key untouched (validate-before-persist still applies:
            // a bad passphrase throws inside importer.import before any keystore write).
            val (parsed, keyHandle) = try {
                val p = importer.import(p12Bytes, passphrase, stagingAlias)
                val k = importer.loadPrivateKey(stagingAlias)
                    ?: error("Imported key not found in AndroidKeyStore slot '$stagingAlias'")
                val metadata = StoredIdentityMetadata(
                    alias = p.alias,
                    keyAlgorithm = p.keyAlgorithm,
                    keyStoreAlias = stagingAlias,
                    certificateChain = p.fullChain,
                )
                certStore.save(metadata) // COMMIT — the single atomic durable live-pointer flip.
                p to k
            } catch (e: Exception) {
                runCatching { importer.remove(stagingAlias) } // best-effort: drop any half-written key
                throw e
            }

            // Committed: publish the new identity (AndroidKeyStore handle, non-exportable) and GC the old.
            liveOverride = Box(
                LiveIdentity(
                    InstalledIdentity(
                        alias = parsed.alias,
                        keyAlgorithm = parsed.keyAlgorithm,
                        privateKey = keyHandle,
                        leafCertificate = parsed.leafCertificate,
                        issuerCertificates = parsed.issuerCertificates,
                    ),
                    stagingAlias,
                ),
            )
            liveAlias?.let { old -> runCatching { importer.remove(old) } } // best-effort GC of old slot
            // Next handshake uses the new identity; drop pooled/resumed connections holding the old one.
            sharedClient.connectionPool.evictAll()
            parsed.summary()
        }

    /** The ping-pong slot to stage into: whichever physical slot is NOT the current live one. */
    private fun stagingSlotFor(liveAlias: String?): String =
        if (liveAlias == importer.primarySlot) importer.secondarySlot else importer.primarySlot

    /**
     * Reconstruct the installed identity from storage at startup, resolving the live key SOLELY via
     * the persisted pointer ([StoredIdentityMetadata.keyStoreAlias]). A storage/decrypt fault must not
     * crash launch — it degrades to "no identity" and is logged (never silently dropped), mirroring
     * iOS `loadedIdentityOrNil`.
     */
    private fun loadInstalledOrNull(): LiveIdentity? {
        return try {
            val metadata = certStore.load() ?: return null
            val keyHandle = importer.loadPrivateKey(metadata.keyStoreAlias) ?: return null
            LiveIdentity(
                InstalledIdentity(
                    alias = metadata.alias,
                    keyAlgorithm = metadata.keyAlgorithm,
                    privateKey = keyHandle,
                    leafCertificate = metadata.leafCertificate,
                    issuerCertificates = metadata.issuerCertificates,
                ),
                metadata.keyStoreAlias,
            )
        } catch (e: Exception) {
            Log.e(TAG, "Failed to load stored device identity; continuing without a client cert", e)
            null
        }
    }

    private fun buildSocketFactory(): SSLSocketFactory {
        val keyManager = ReReadingX509KeyManager { currentLive()?.installed }
        val context = SSLContext.getInstance("TLS")
        context.init(arrayOf<KeyManager>(keyManager), arrayOf<TrustManager>(systemTrustManager), null)
        return context.socketFactory
    }

    private fun systemDefaultTrustManager(): X509TrustManager {
        val factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
        factory.init(null as KeyStore?)
        return factory.trustManagers.filterIsInstance<X509TrustManager>().firstOrNull()
            ?: error("No system X509TrustManager available")
    }

    private companion object {
        private const val TAG = "IdentityRepository"
    }
}
