package com.osglab.account.features.oobe

import com.nimbusds.jwt.SignedJWT
import com.nimbusds.jwt.JWTClaimsSet
import com.nimbusds.jose.JWSAlgorithm
import com.nimbusds.jose.JWSHeader
import com.nimbusds.jose.crypto.MACSigner
import com.osglab.account.config.AppleServiceEnvironment
import com.osglab.account.config.IntegrityConfig
import com.osglab.account.config.IntegrityPolicy
import com.osglab.account.features.gateway.models.GatewaySubjectType
import com.osglab.account.features.gateway.models.ProviderUsage
import com.osglab.account.features.integrity.AppAttestChallenge
import com.osglab.account.features.integrity.AppAttestChallengePurpose
import com.osglab.account.features.integrity.AppAttestCrypto
import com.osglab.account.features.integrity.AppAttestKeyStatus
import com.osglab.account.features.integrity.AppAttestRepository
import com.osglab.account.features.integrity.AppAttestService
import com.osglab.account.features.integrity.AppAttestRejectedException
import com.osglab.account.features.integrity.AttestedKeyMaterial
import com.osglab.account.features.integrity.ConsumedChallenge
import com.osglab.account.features.integrity.StoredAppAttestKey
import io.kotest.assertions.throwables.shouldThrow
import io.kotest.core.spec.style.StringSpec
import io.kotest.matchers.collections.shouldContainExactlyInAnyOrder
import io.kotest.matchers.shouldBe
import java.security.MessageDigest
import java.time.Clock
import java.time.Duration
import java.time.Instant
import java.time.ZoneId
import java.util.Base64
import java.util.UUID

class OobeGrantServiceTest : StringSpec({
    "canonical assertion is server-owned and binds all fixed permissions" {
        val challenge = ByteArray(32) { it.toByte() }
        val payload = OobeContract.canonicalAssertionPayload(challenge, KEY_ID, INSTALLATION_ID).decodeToString()

        payload shouldBe """
            osg-app-attest-v1
            purpose=oobe-gateway-grant
            challenge=AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8
            key_id=$KEY_ID
            installation_id=$INSTALLATION_ID
            scopes=ai,polish
            features=ask_ai,clipboard_reply,clipboard_translate,voice_input
            grant_ttl_seconds=1800
            access_ttl_seconds=300

        """.trimIndent()
    }

    "issues a distinct short-lived OOBE token after the canonical assertion" {
        val clock = MutableClock(NOW)
        val challenge = ByteArray(32) { 7 }
        val expectedHash = sha256(
            OobeContract.canonicalAssertionPayload(challenge, KEY_ID, INSTALLATION_ID),
        )
        val repository = FakeOobeRepository()
        val service = service(repository, expectedHash, clock)

        val tokens = service.create(request(challenge))
        val jwt = SignedJWT.parse(tokens.accessToken).jwtClaimsSet

        jwt.getStringClaim("typ") shouldBe "oobe_gateway_access"
        jwt.subject.startsWith("oobe:") shouldBe true
        jwt.getStringListClaim("scp").shouldContainExactlyInAnyOrder("polish", "ai")
        jwt.getStringListClaim("features").shouldContainExactlyInAnyOrder(
            "voice_input",
            "clipboard_translate",
            "clipboard_reply",
            "ask_ai",
        )
        Duration.between(jwt.issueTime.toInstant(), jwt.expirationTime.toInstant()) shouldBe
            Duration.ofMinutes(5)
        Duration.between(NOW, Instant.parse(tokens.refreshExpiresAt)) shouldBe Duration.ofMinutes(30)
        service.authenticate(tokens.accessToken)?.subjectType shouldBe GatewaySubjectType.OOBE

        val overScoped = SignedJWT(
            JWSHeader(JWSAlgorithm.HS256),
            JWTClaimsSet.Builder(jwt)
                .claim("scp", listOf("ai", "agent", "polish"))
                .build(),
        ).apply { sign(MACSigner(ByteArray(32) { 1 })) }.serialize()
        service.authenticate(overScoped) shouldBe null

        val accountTyped = SignedJWT(
            JWSHeader(JWSAlgorithm.HS256),
            JWTClaimsSet.Builder(jwt)
                .claim("typ", "gateway_access")
                .build(),
        ).apply { sign(MACSigner(ByteArray(32) { 1 })) }.serialize()
        service.authenticate(accountTyped) shouldBe null

        clock.now = NOW.plus(Duration.ofMinutes(5))
        service.authenticate(tokens.accessToken) shouldBe null
    }

    "rejects an assertion generated for a different installation payload" {
        val challenge = ByteArray(32) { 9 }
        val signedHash = sha256(
            OobeContract.canonicalAssertionPayload(challenge, KEY_ID, INSTALLATION_ID),
        )
        val service = service(FakeOobeRepository(), signedHash, MutableClock(NOW))

        shouldThrow<AppAttestRejectedException> {
            service.create(request(challenge).copy(installationId = UUID.randomUUID().toString()))
        }
    }

    "refresh cannot extend the original grant TTL" {
        val clock = MutableClock(NOW)
        val challenge = ByteArray(32) { 5 }
        val expectedHash = sha256(
            OobeContract.canonicalAssertionPayload(challenge, KEY_ID, INSTALLATION_ID),
        )
        val service = service(FakeOobeRepository(), expectedHash, clock)
        val created = service.create(request(challenge))

        clock.now = NOW.plus(Duration.ofMinutes(29))
        val refreshed = service.refresh(created.refreshToken, "oobe-refresh-1")

        refreshed.refreshExpiresAt shouldBe created.refreshExpiresAt
        refreshed.accessExpiresAt shouldBe NOW.plus(Duration.ofMinutes(30)).toString()
    }
})

private fun service(
    repository: OobeRepository,
    expectedHash: ByteArray,
    clock: Clock,
): OobeGrantService {
    val appAttestRepository = FakeAppAttestRepository()
    val appAttest = AppAttestService(
        repository = appAttestRepository,
        crypto = HashCheckingAppAttestCrypto(expectedHash),
        config = IntegrityConfig(
            deviceCheckPolicy = IntegrityPolicy.ENFORCE,
            appAttestPolicy = IntegrityPolicy.ENFORCE,
            appleEnvironment = AppleServiceEnvironment.PRODUCTION,
        ),
        clock = clock,
    )
    return OobeGrantService(
        repository = repository,
        appAttest = appAttest,
        settings = OobeTokenSettings(
            issuer = "osg-test",
            audience = "osg-gateway-test",
            accessTokenHmacSecret = ByteArray(32) { 1 },
            refreshTokenHmacSecret = ByteArray(32) { 2 },
        ),
        clock = clock,
    )
}

private class HashCheckingAppAttestCrypto(
    private val expectedHash: ByteArray,
) : AppAttestCrypto {
    override suspend fun validateAttestation(
        attestationObject: ByteArray,
        keyId: String,
        challenge: ByteArray,
    ): AttestedKeyMaterial = error("not used")

    override suspend fun validateAssertion(
        assertionObject: ByteArray,
        clientDataHash: ByteArray,
        publicKey: ByteArray,
        lastCounter: Long,
    ): Long {
        if (!MessageDigest.isEqual(clientDataHash, expectedHash)) {
            throw AppAttestRejectedException("canonical payload mismatch")
        }
        return lastCounter + 1
    }
}

private class FakeAppAttestRepository : AppAttestRepository {
    private var counter = 0L

    override suspend fun createChallenge(challenge: AppAttestChallenge) = Unit

    override suspend fun consumeChallenge(
        id: UUID,
        purpose: AppAttestChallengePurpose,
        keyId: String,
        challengeHash: String,
        accountId: UUID?,
        now: Instant,
    ): ConsumedChallenge = ConsumedChallenge.Valid

    override suspend fun saveKey(key: StoredAppAttestKey): Boolean = true

    override suspend fun findKey(keyId: String): StoredAppAttestKey =
        StoredAppAttestKey(
            keyId = keyId,
            publicKey = byteArrayOf(1),
            receipt = byteArrayOf(1),
            counter = counter,
            accountId = null,
            status = AppAttestKeyStatus.ACTIVE,
        )

    override suspend fun updateCounter(
        keyId: String,
        expectedCounter: Long,
        newCounter: Long,
        now: Instant,
    ): Boolean {
        if (expectedCounter != counter || newCounter <= counter) return false
        counter = newCounter
        return true
    }

    override suspend fun bindKeyToAccount(keyId: String, accountId: UUID, now: Instant): Boolean = true
}

private class FakeOobeRepository : OobeRepository {
    private val subjects = mutableMapOf<Pair<String, String>, OobeSubject>()
    private val grants = mutableMapOf<String, OobeGrant>()
    private val refreshes = mutableMapOf<String, StoredOobeRefresh>()

    override suspend fun findOrCreateSubject(
        keyId: String,
        installationHash: String,
        subjectId: String,
        now: Instant,
    ): OobeSubject = subjects.getOrPut(keyId to installationHash) {
        OobeSubject(subjectId, keyId, installationHash)
    }

    override suspend fun createGrant(grant: NewOobeGrant, now: Instant): StoredOobeRefresh {
        grants[grant.grant.id] = grant.grant
        return StoredOobeRefresh(
            grant.grant,
            grant.refreshTokenId,
            grant.refreshFamilyId,
            grant.refreshExpiresAt,
        ).also { refreshes[grant.refreshTokenHash] = it }
    }

    override suspend fun rotateRefresh(
        currentTokenHash: String,
        rotationIdempotencyKey: String,
        newTokenId: String,
        newTokenHash: String,
        newExpiresAt: Instant,
        now: Instant,
    ): OobeRefreshRotationResult {
        val current = refreshes[currentTokenHash] ?: return OobeRefreshRotationResult.Invalid
        if (!current.expiresAt.isAfter(now) || !current.grant.expiresAt.isAfter(now)) {
            return OobeRefreshRotationResult.Invalid
        }
        val replacement = StoredOobeRefresh(
            grant = current.grant,
            tokenId = newTokenId,
            familyId = current.familyId,
            expiresAt = minOf(newExpiresAt, current.grant.expiresAt),
        )
        refreshes[newTokenHash] = replacement
        return OobeRefreshRotationResult.Rotated(replacement)
    }

    override suspend fun findActiveGrant(grantId: String, subjectId: String, now: Instant): OobeGrant? =
        grants[grantId]?.takeIf { it.subjectId == subjectId && it.expiresAt.isAfter(now) }

    override suspend fun claim(
        request: OobeProviderRequest,
        expiresAt: Instant,
        now: Instant,
    ): OobeRequestClaim? = error("not used")

    override suspend fun markStarted(claim: OobeRequestClaim) = error("not used")
    override suspend fun consume(claim: OobeRequestClaim, usage: ProviderUsage) = error("not used")
    override suspend fun release(claim: OobeRequestClaim, errorCode: String) = error("not used")
    override suspend fun markManualReview(claim: OobeRequestClaim, errorCode: String) = error("not used")
}

private class MutableClock(var now: Instant) : Clock() {
    override fun getZone(): ZoneId = ZoneId.of("UTC")
    override fun withZone(zone: ZoneId): Clock = this
    override fun instant(): Instant = now
}

private fun request(challenge: ByteArray) = CreateOobeGrantRequest(
    challengeId = UUID.randomUUID().toString(),
    challenge = Base64.getUrlEncoder().withoutPadding().encodeToString(challenge),
    keyId = KEY_ID,
    installationId = INSTALLATION_ID,
    assertion = Base64.getEncoder().encodeToString(byteArrayOf(1)),
)

private fun sha256(value: ByteArray): ByteArray = MessageDigest.getInstance("SHA-256").digest(value)

private val NOW = Instant.parse("2026-08-21T00:00:00Z")
private val INSTALLATION_ID = UUID.fromString("10000000-0000-0000-0000-000000000001").toString()
private val KEY_ID = Base64.getEncoder().encodeToString(ByteArray(32) { 3 })