package com.rainy.token.data.proxy

import com.rainy.token.data.repository.CredentialRepository
import com.rainy.token.domain.model.Credential
import com.rainy.token.domain.service.ServiceType
import javax.inject.Inject
import javax.inject.Singleton
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

/**
 * 账号池：从 CredentialRepository 实时取号，负责轮询与会话粘性。
 *
 * - 多账号服务（Trae / WorkBuddy / Sub2API）按 ProviderKind 分组；
 * - 会话粘性：同一 conversation 头持续命中同一账号，避免上下文错乱；
 * - 区域匹配：优先选 credential.region 与路由 region 一致的账号，找不到时退回全部账号。
 */
@Singleton
class AccountPool @Inject constructor(
    private val credentialRepository: CredentialRepository,
) {

    data class PooledAccount(
        val kind: ProviderKind,
        val region: ProviderRegion? = null,
        val accountId: String? = null,
        val label: String? = null,
    )

    private val mutex = Mutex()
    private val sessionSticky = mutableMapOf<String, String>()
    private val cursor = mutableMapOf<ProviderKind, Int>()

    suspend fun next(
        kind: ProviderKind,
        region: ProviderRegion? = null,
        sessionKey: String? = null,
    ): PooledAccount? = mutex.withLock {
        val service = serviceFor(kind) ?: return null
        val accounts = credentialRepository.accountsFor(service)
        val allIds = accounts.map { it.id }
        val matchedIds = if (region == null) {
            allIds
        } else {
            allIds.filter { accountId -> regionMatches(service, accountId, region) }
                .ifEmpty { allIds }
        }
        if (matchedIds.isEmpty()) return null

        val accountId = if (sessionKey != null) {
            sessionSticky[sessionKey]
                ?.takeIf { it in matchedIds }
                ?: pickRoundRobin(kind, matchedIds).also { sessionSticky[sessionKey] = it }
        } else {
            pickRoundRobin(kind, matchedIds)
        }

        PooledAccount(
            kind = kind,
            region = region,
            accountId = accountId,
            label = accounts.firstOrNull { it.id == accountId }?.label,
        )
    }

    private suspend fun regionMatches(service: ServiceType, accountId: String, region: ProviderRegion): Boolean {
        val credential = credentialRepository.get(service, accountId) ?: return false
        return when (credential) {
            is Credential.TraeCredential -> credential.region.equals(region.name, ignoreCase = true)
            is Credential.WorkBuddyCredential -> credential.region.equals(region.name, ignoreCase = true)
            else -> true
        }
    }

    private fun pickRoundRobin(kind: ProviderKind, ids: List<String>): String {
        val idx = cursor[kind] ?: 0
        cursor[kind] = (idx + 1) % ids.size
        return ids[idx]
    }

    private fun serviceFor(kind: ProviderKind): ServiceType? = when (kind) {
        ProviderKind.WORKBUDDY_CN, ProviderKind.WORKBUDDY_INTL -> ServiceType.WORKBUDDY
        ProviderKind.TRAE_CN, ProviderKind.TRAE_INTL -> ServiceType.TRAE
        ProviderKind.OPENAI_COMPATIBLE -> ServiceType.SUB2API
        else -> null
    }
}