package com.correx.testing.contracts.model.inference

import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.StageId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.DefaultInferenceRouter
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoEligibleProviderException
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.ProviderRegistry
import com.correx.core.inference.RoutingStrategy
import com.correx.testing.fixtures.inference.MockInferenceProvider
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertSame
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows
import kotlin.time.Duration.Companion.milliseconds

class DefaultInferenceRouterTest {

    private val stage = StageId("s1")

    private fun provider(id: String, vararg caps: ModelCapability): MockInferenceProvider =
        MockInferenceProvider(
            id = ProviderId(id),
            declaredCapabilities = caps.map { CapabilityScore(it, 1.0) }.toSet(),
        )

    private fun registryOf(vararg providers: InferenceProvider): ProviderRegistry =
        object : ProviderRegistry {
            override fun register(provider: InferenceProvider) = Unit
            override fun resolve(capability: ModelCapability) =
                providers.filter { p -> p.capabilities().any { it.capability == capability } }

            override fun listAll() = providers.toList()
            override suspend fun healthCheckAll() = providers.associate { it.id to ProviderHealth.Healthy }
        }

    private fun firstStrategy(): RoutingStrategy =
        RoutingStrategy { candidates, _ ->
            candidates.firstOrNull() ?: throw NoEligibleProviderException(StageId("?"), emptySet())
        }

    private fun throwingStrategy(): RoutingStrategy =
        RoutingStrategy { _, requiredCapabilities ->
            throw NoEligibleProviderException(
                StageId("?"),
                requiredCapabilities,
            )
        }

    // ── capability matching ───────────────────────────────────────────────────

    @Test
    fun `routes to provider that matches required capability`(): Unit = runBlocking {
        val p = provider("a", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(p), firstStrategy())
        assertSame(p, router.route(stage, setOf(ModelCapability.General)))
    }

    @Test
    fun `deduplicates providers that satisfy multiple required capabilities`(): Unit = runBlocking {
        val p = provider("a", ModelCapability.General, ModelCapability.Coding)
        val selected = mutableListOf<InferenceProvider>()
        val captureStrategy = RoutingStrategy { candidates, _ ->
            selected.addAll(candidates)
            candidates.first()
        }
        val router = DefaultInferenceRouter(
            registryOf(p),
            captureStrategy,
        )
        router.route(stage, setOf(ModelCapability.General, ModelCapability.Coding))
        assert(selected.count { it.id == p.id } == 1) { "provider appeared ${selected.size} times, expected 1" }
    }

    // ── empty capabilities falls back to all providers ────────────────────────

    @Test
    fun `empty required capabilities routes from all registered providers`(): Unit = runBlocking {
        val p1 = provider("a", ModelCapability.General)
        val p2 = provider("b", ModelCapability.Coding)
        val router = DefaultInferenceRouter(registryOf(p1, p2), firstStrategy())
        val result = router.route(stage, emptySet())
        assertSame(p1, result)
    }

    // ── no candidates throws ──────────────────────────────────────────────────

    @Test
    fun `throws NoEligibleProviderException when no provider satisfies capability`() {
        val router = DefaultInferenceRouter(registryOf(), throwingStrategy())
        assertThrows<NoEligibleProviderException> {
            runBlocking { router.route(stage, setOf(ModelCapability.General)) }
        }
    }

    @Test
    fun `throws NoEligibleProviderException when registry is empty and capabilities are empty`() {
        val router = DefaultInferenceRouter(registryOf(), throwingStrategy())
        assertThrows<NoEligibleProviderException> {
            runBlocking { router.route(stage, emptySet()) }
        }
    }

    // ── unavailable providers are filtered ───────────────────────────────────

    @Test
    fun `filters out unavailable providers before selection`(): Unit = runBlocking {
        val p1 = MockInferenceProvider(
            id = com.correx.core.events.types.ProviderId("a"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
            health = ProviderHealth.Unavailable("down"),
        )
        val p2 = provider("b", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(p1, p2), firstStrategy())
        val result = router.route(stage, setOf(ModelCapability.General))
        assertSame(p2, result)
    }

    // ── 3-arg route with modelId ───────────────────────────────────────────

    @Test
    fun `3-arg route selects provider by exact id match`(): Unit = runBlocking {
        val p1 = provider("llama-cpp:fast-model", ModelCapability.General)
        val p2 = provider("llama-cpp:big-model", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(p1, p2), firstStrategy())
        val result = router.route(stage, setOf(ModelCapability.General), "llama-cpp:fast-model")
        assertSame(p1, result)
    }

    @Test
    fun `3-arg route falls back to capability routing when modelId matches no provider`(): Unit = runBlocking {
        val p1 = provider("llama-cpp:only-one", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(p1), firstStrategy())
        val result = router.route(stage, setOf(ModelCapability.General), "nonexistent-model")
        assertSame(p1, result)
    }

    @Test
    fun `3-arg route with null modelId delegates to 2-arg form`(): Unit = runBlocking {
        val p = provider("llama-cpp:model", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(p), firstStrategy())
        val result = router.route(stage, setOf(ModelCapability.General), null)
        assertSame(p, result)
    }

    @Test
    fun `3-arg route falls back to capability routing when matched provider is unhealthy`(): Unit = runBlocking {
        val sick = MockInferenceProvider(
            id = ProviderId("llama-cpp:sick"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
            health = ProviderHealth.Unavailable("down"),
        )
        val healthy = provider("llama-cpp:healthy", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(sick, healthy), firstStrategy())
        val result = router.route(stage, setOf(ModelCapability.General), "llama-cpp:sick")
        assertSame(healthy, result)
    }

    // ── bounded wait for a briefly-absent provider (#299) ─────────────────────

    @Test
    fun `waits for the sole capable provider to recover instead of failing immediately`(): Unit = runBlocking {
        var healthChecks = 0
        val recovering = object : InferenceProvider by provider("a", ModelCapability.ToolCalling) {
            override suspend fun healthCheck(): ProviderHealth {
                healthChecks++
                return if (healthChecks < 3) ProviderHealth.Unavailable("connection dropped") else ProviderHealth.Healthy
            }
        }
        val router = DefaultInferenceRouter(
            registryOf(recovering),
            firstStrategy(),
            unavailableRetryAttempts = 5,
            unavailableRetryDelay = 5.milliseconds,
        )
        val result = router.route(stage, setOf(ModelCapability.ToolCalling))
        assertSame(recovering, result)
        assertTrue(healthChecks >= 3) { "expected at least 3 health checks, got $healthChecks" }
    }

    @Test
    fun `still throws NoEligibleProviderException if the sole provider never recovers within the bound`() {
        val neverRecovers = MockInferenceProvider(
            id = ProviderId("a"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.ToolCalling, 1.0)),
            health = ProviderHealth.Unavailable("still down"),
        )
        val router = DefaultInferenceRouter(
            registryOf(neverRecovers),
            throwingStrategy(),
            unavailableRetryAttempts = 2,
            unavailableRetryDelay = 5.milliseconds,
        )
        assertThrows<NoEligibleProviderException> {
            runBlocking { router.route(stage, setOf(ModelCapability.ToolCalling)) }
        }
    }

    @Test
    fun `fails fast without waiting when the capability was never configured on any provider`(): Unit = runBlocking {
        val p = provider("a", ModelCapability.General) // does not declare ToolCalling
        val router = DefaultInferenceRouter(
            registryOf(p),
            throwingStrategy(),
            unavailableRetryAttempts = 5,
            unavailableRetryDelay = 10_000.milliseconds, // would time the test out if the wait loop ran
        )
        assertThrows<NoEligibleProviderException> {
            router.route(stage, setOf(ModelCapability.ToolCalling))
        }
    }

    // ── event-driven health gating (#300) ─────────────────────────────────────

    @Test
    fun `reportFailure gates the very next route call without waiting for a health poll`(): Unit = runBlocking {
        var healthCheckCount = 0
        val alwaysClaimsHealthy = object : InferenceProvider by provider("a", ModelCapability.General) {
            override suspend fun healthCheck(): ProviderHealth {
                healthCheckCount++
                return ProviderHealth.Healthy // the provider's own health probe hasn't caught up yet
            }
        }
        val backup = provider("b", ModelCapability.General)
        val router = DefaultInferenceRouter(registryOf(alwaysClaimsHealthy, backup), firstStrategy())

        // Sanity: before reportFailure, routes to the first (still "healthy") provider.
        assertSame(alwaysClaimsHealthy, router.route(stage, setOf(ModelCapability.General)))

        router.reportFailure(alwaysClaimsHealthy.id, "connection reset")

        // Immediately after — no delay, no waiting for the next poll — routing must avoid it.
        val result = router.route(stage, setOf(ModelCapability.General))
        assertSame(backup, result)
    }
}
