package com.osglab.account.features.oobe

import com.osglab.account.features.gateway.models.GatewayCapability
import com.osglab.account.features.gateway.models.GatewayRequestPurpose
import com.osglab.account.features.gateway.models.GatewayTaskKind
import com.osglab.account.features.gateway.models.OobeFeature
import kotlinx.serialization.Serializable
import java.time.Instant
import java.util.Base64

@Serializable
data class CreateOobeGrantRequest(
    val challengeId: String,
    val challenge: String,
    val keyId: String,
    val installationId: String,
    val assertion: String,
)

@Serializable
data class RefreshOobeGrantRequest(val refreshToken: String)

@Serializable
data class OobeGrantTokens(
    val grantId: String,
    val scopes: Set<GatewayCapability>,
    val features: Set<OobeFeature>,
    val accessToken: String,
    val accessExpiresAt: String,
    val refreshToken: String,
    val refreshExpiresAt: String,
)

data class OobeFeaturePolicy(
    val capability: GatewayCapability,
    val taskKind: GatewayTaskKind,
)

object OobeContract {
    val scopes: Set<GatewayCapability> = setOf(GatewayCapability.POLISH, GatewayCapability.AI)
    val features: Set<OobeFeature> = OobeFeature.entries.toSet()

    fun policy(feature: OobeFeature): OobeFeaturePolicy = when (feature) {
        OobeFeature.VOICE_INPUT ->
            OobeFeaturePolicy(GatewayCapability.POLISH, GatewayTaskKind.DICTATION_POLISH)
        OobeFeature.CLIPBOARD_TRANSLATE,
        OobeFeature.CLIPBOARD_REPLY ->
            OobeFeaturePolicy(GatewayCapability.AI, GatewayTaskKind.CLIPBOARD_TRANSFORM)
        OobeFeature.ASK_AI ->
            OobeFeaturePolicy(GatewayCapability.AI, GatewayTaskKind.AI_QUESTION)
    }

    fun canonicalAssertionPayload(
        challenge: ByteArray,
        keyId: String,
        installationId: String,
    ): ByteArray = buildString {
        appendLine("osg-app-attest-v1")
        appendLine("purpose=oobe-gateway-grant")
        appendLine("challenge=${BASE64_URL.encodeToString(challenge)}")
        appendLine("key_id=$keyId")
        appendLine("installation_id=$installationId")
        appendLine("scopes=ai,polish")
        appendLine("features=ask_ai,clipboard_reply,clipboard_translate,voice_input")
        appendLine("grant_ttl_seconds=1800")
        appendLine("access_ttl_seconds=300")
    }.toByteArray(Charsets.UTF_8)
}

data class OobeSubject(
    val id: String,
    val keyId: String,
    val installationHash: String,
)

data class OobeGrant(
    val id: String,
    val subjectId: String,
    val expiresAt: Instant,
    val revokedAt: Instant? = null,
)

data class NewOobeGrant(
    val grant: OobeGrant,
    val refreshTokenId: String,
    val refreshFamilyId: String,
    val refreshTokenHash: String,
    val refreshExpiresAt: Instant,
)

data class StoredOobeRefresh(
    val grant: OobeGrant,
    val tokenId: String,
    val familyId: String,
    val expiresAt: Instant,
)

sealed interface OobeRefreshRotationResult {
    data class Rotated(val refresh: StoredOobeRefresh) : OobeRefreshRotationResult
    data object Invalid : OobeRefreshRotationResult
    data object ReuseDetected : OobeRefreshRotationResult
}

data class OobeRequestClaim(
    val subjectId: String,
    val grantId: String,
    val feature: OobeFeature,
    val requestId: String,
)

data class OobeProviderRequest(
    val subjectId: String,
    val grantId: String,
    val feature: OobeFeature,
    val requestId: String,
    val providerId: String,
    val capability: GatewayCapability,
    val purpose: GatewayRequestPurpose,
)

enum class OobeProviderRequestState {
    CLAIMED,
    STARTED,
    SUCCEEDED,
    RELEASED,
    MANUAL_REVIEW,
}

private val BASE64_URL: Base64.Encoder = Base64.getUrlEncoder().withoutPadding()