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.NoEligibleProviderException
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 org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows

class CapabilityAwareRoutingStrategyTest {

    private val strategy = CapabilityAwareRoutingStrategy()

    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("unused")
            override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
            override fun capabilities(): Set<CapabilityScore> =
                caps.map { CapabilityScore(it.first, it.second) }.toSet()
        }

    @Test
    fun `picks the highest-scoring provider for the required capability`() {
        val weak = fakeProvider("weak", ModelCapability.ToolCalling to 0.4, ModelCapability.General to 1.0)
        val strong = fakeProvider("strong", ModelCapability.ToolCalling to 0.95, ModelCapability.General to 1.0)
        val selected = strategy.select(listOf(weak, strong), setOf(ModelCapability.ToolCalling))
        assertEquals("strong", selected.id.value)
    }

    @Test
    fun `still hard-filters to providers declaring the required capability`() {
        val noTools = fakeProvider("noTools", ModelCapability.General to 1.0)
        val tools = fakeProvider("tools", ModelCapability.ToolCalling to 0.3)
        val selected = strategy.select(listOf(noTools, tools), setOf(ModelCapability.ToolCalling))
        assertEquals("tools", selected.id.value)
    }

    @Test
    fun `throws when no candidate declares the required capability`() {
        val general = fakeProvider("general", ModelCapability.General to 1.0)
        assertThrows<NoEligibleProviderException> {
            strategy.select(listOf(general), setOf(ModelCapability.ToolCalling))
        }
    }

    @Test
    fun `keeps list order on a score tie (first-available parity)`() {
        val a = fakeProvider("a", ModelCapability.ToolCalling to 0.5)
        val b = fakeProvider("b", ModelCapability.ToolCalling to 0.5)
        assertEquals("a", strategy.select(listOf(a, b), setOf(ModelCapability.ToolCalling)).id.value)
    }

    @Test
    fun `with no required capabilities ranks by total declared score`() {
        val lean = fakeProvider("lean", ModelCapability.General to 0.5)
        val rich = fakeProvider("rich", ModelCapability.General to 1.0, ModelCapability.Coding to 0.9)
        assertEquals("rich", strategy.select(listOf(lean, rich), emptySet()).id.value)
    }
}
