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.assertEquals
import org.junit.jupiter.api.Assertions.assertSame
import org.junit.jupiter.api.Test
import kotlin.time.Duration.Companion.seconds
import kotlin.time.TestTimeSource

class DefaultInferenceRouterCacheTest {

    private val stage = StageId("cache-test")

    private class CountingProvider(
        delegate: MockInferenceProvider,
    ) : InferenceProvider by delegate {
        var healthCheckCount: Int = 0
            private set
        private var _health: ProviderHealth = ProviderHealth.Healthy

        fun setHealth(h: ProviderHealth) { _health = h }

        override suspend fun healthCheck(): ProviderHealth {
            healthCheckCount++
            return _health
        }
    }

    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, required ->
            candidates.firstOrNull() ?: throw NoEligibleProviderException(StageId("?"), required)
        }

    // ── (a) repeated route() within TTL invokes health check once per provider ──

    @Test
    fun `health check is cached within TTL`(): Unit = runBlocking {
        val timeSource = TestTimeSource()
        val mock = MockInferenceProvider(
            id = ProviderId("p1"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
        )
        val counting = CountingProvider(mock)
        val router = DefaultInferenceRouter(
            registry = registryOf(counting),
            strategy = firstStrategy(),
            cacheTtl = 5.seconds,
            timeSource = timeSource,
        )

        router.route(stage, setOf(ModelCapability.General))
        router.route(stage, setOf(ModelCapability.General))
        router.route(stage, setOf(ModelCapability.General))

        // Cache hit for all subsequent calls: health checked once for the filter pass,
        // plus the post-selection re-check on the first call (uncached path), then cached.
        // After the first route() the entry is cached; the two subsequent calls use the cache.
        // Post-selection re-check always goes live, so total = 1 (filter) + 1 (post-select)
        // on first call, then 0 (filter cached) + 1 (post-select live) * 2 more = 4 total.
        // But the cache only governs the filter pass; post-select is always live by design.
        // So: first call = 2 checks (filter + post-select), subsequent calls = 1 each (post-select).
        assertEquals(4, counting.healthCheckCount)
    }

    // ── (b) after TTL expires, health is re-checked ──────────────────────────

    @Test
    fun `health check is refreshed after TTL expires`(): Unit = runBlocking {
        val timeSource = TestTimeSource()
        val mock = MockInferenceProvider(
            id = ProviderId("p1"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
        )
        val counting = CountingProvider(mock)
        val router = DefaultInferenceRouter(
            registry = registryOf(counting),
            strategy = firstStrategy(),
            cacheTtl = 5.seconds,
            timeSource = timeSource,
        )

        router.route(stage, setOf(ModelCapability.General))
        val countAfterFirstCall = counting.healthCheckCount

        // Advance past TTL
        timeSource.plusAssign(6.seconds)

        router.route(stage, setOf(ModelCapability.General))

        // After TTL expiry a new filter-pass health check must fire
        assert(counting.healthCheckCount > countAfterFirstCall + 1) {
            "Expected more than ${countAfterFirstCall + 1} total checks, got ${counting.healthCheckCount}"
        }
    }

    // ── (c) unhealthy cached entry causes provider skip ───────────────────────

    @Test
    fun `cached unavailable entry causes provider to be skipped`(): Unit = runBlocking {
        val timeSource = TestTimeSource()
        val mock1 = MockInferenceProvider(
            id = ProviderId("p1"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
        )
        val mock2 = MockInferenceProvider(
            id = ProviderId("p2"),
            declaredCapabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
        )
        val counting1 = CountingProvider(mock1)
        val counting2 = CountingProvider(mock2)
        val router = DefaultInferenceRouter(
            registry = registryOf(counting1, counting2),
            strategy = firstStrategy(),
            cacheTtl = 5.seconds,
            timeSource = timeSource,
        )

        // Warm the cache with p1 healthy
        router.route(stage, setOf(ModelCapability.General))

        // Mark p1 unavailable in the counting provider — but cache still says healthy
        counting1.setHealth(ProviderHealth.Unavailable("went down"))

        // Advance past TTL so cache refreshes
        timeSource.plusAssign(6.seconds)

        val selected = router.route(stage, setOf(ModelCapability.General))
        assertSame(counting2, selected)
    }
}
