package com.correx.infrastructure.inference

import com.correx.core.events.types.ProviderId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.Token
import com.correx.core.inference.Tokenizer
import com.correx.core.utils.TypeId
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class DefaultProviderRegistryTest {

    private fun fakeProvider(
        id: String,
        vararg caps: Pair<ModelCapability, Double>,
    ): InferenceProvider = object : InferenceProvider {
        override val id: ProviderId = TypeId(id)
        override val name: String = id
        override val tokenizer: Tokenizer = object : Tokenizer {
            override suspend fun tokenize(text: String): List<Token> = emptyList()
            override suspend fun countTokens(text: String): Int = 0
        }
        override suspend fun infer(request: InferenceRequest): InferenceResponse =
            error("not implemented in test")
        override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
        override fun capabilities(): Set<CapabilityScore> =
            caps.map { (cap, score) -> CapabilityScore(cap, score) }.toSet()
    }

    @Test
    fun `listAll returns all registered providers`() {
        val p1 = fakeProvider("p1", ModelCapability.Coding to 0.9)
        val p2 = fakeProvider("p2", ModelCapability.General to 0.8)
        val registry = DefaultProviderRegistry(listOf(p1, p2))
        assertEquals(2, registry.listAll().size)
    }

    @Test
    fun `register adds provider`() {
        val registry = DefaultProviderRegistry()
        val p = fakeProvider("p1", ModelCapability.Coding to 0.9)
        registry.register(p)
        assertEquals(1, registry.listAll().size)
    }

    @Test
    fun `resolve returns providers with matching capability ordered by score desc`() {
        val low = fakeProvider("low", ModelCapability.Coding to 0.5)
        val high = fakeProvider("high", ModelCapability.Coding to 0.9)
        val other = fakeProvider("other", ModelCapability.General to 0.8)
        val registry = DefaultProviderRegistry(listOf(low, high, other))

        val resolved = registry.resolve(ModelCapability.Coding)
        assertEquals(2, resolved.size)
        assertEquals("high", resolved.first().id.value)
    }

    @Test
    fun `resolve returns empty list when no provider matches`() {
        val registry = DefaultProviderRegistry(listOf(fakeProvider("p", ModelCapability.General to 1.0)))
        assertTrue(registry.resolve(ModelCapability.Reasoning).isEmpty())
    }

    @Test
    fun `healthCheckAll returns health for all providers`() = runTest {
        val p1 = fakeProvider("p1", ModelCapability.Coding to 1.0)
        val p2 = fakeProvider("p2", ModelCapability.General to 1.0)
        val registry = DefaultProviderRegistry(listOf(p1, p2))
        val health = registry.healthCheckAll()
        assertEquals(2, health.size)
        assertTrue(health.values.all { it is ProviderHealth.Healthy })
    }
}
