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.InferenceProvider
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoEligibleProviderException
import com.correx.core.inference.ProviderHealth
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ModelManager
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.testing.fixtures.inference.MockInferenceProvider
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows

class ManagedInferenceRouterTest {

    private val stage = StageId("s1")

    /** Records every (idempotent) load/ensureLoaded so we can assert which model was made resident. */
    private class FakeModelManager : ModelManager {
        val loads = mutableListOf<String>()
        private var current: ModelDescriptor? = null

        override suspend fun load(descriptor: ModelDescriptor): InferenceProvider {
            loads.add(descriptor.modelId)
            current = descriptor
            return MockInferenceProvider(id = ProviderId("managed:${descriptor.modelId}"))
        }

        override suspend fun unload(modelId: String) { current = null }
        override fun currentModel(): ModelDescriptor? = current
        override fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
    }

    private fun desc(id: String, vararg caps: Pair<ModelCapability, Double>): ModelDescriptor =
        ModelDescriptor(
            modelId = id,
            modelPath = "/models/$id",
            residencyMode = ResidencyMode.PERSISTENT,
            capabilities = caps.map { CapabilityScore(it.first, it.second) }.toSet(),
        )

    @Test
    fun `stage modelId takes precedence over capability and default`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)),
            defaultModelId = "default-model",
        )

        val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = "coder")

        assertEquals("managed:coder", provider.id.value)
        assertEquals(listOf("coder"), manager.loads)
    }

    @Test
    fun `capability match selects a model that covers the required capabilities`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(
                desc("generalist", ModelCapability.General to 1.0),
                desc("coder", ModelCapability.Coding to 1.0),
            ),
            defaultModelId = "generalist",
        )

        val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null)

        assertEquals("managed:coder", provider.id.value)
        assertEquals(listOf("coder"), manager.loads)
    }

    @Test
    fun `capability match prefers the highest summed score`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(
                desc("weak", ModelCapability.Coding to 0.3),
                desc("strong", ModelCapability.Coding to 0.9),
            ),
            defaultModelId = "weak",
        )

        val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null)

        assertEquals("managed:strong", provider.id.value)
    }

    @Test
    fun `falls back to default when no modelId and no required capabilities`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)),
            defaultModelId = "default-model",
        )

        val provider = router.route(stage, emptySet(), modelId = null)

        assertEquals("managed:default-model", provider.id.value)
        assertEquals(listOf("default-model"), manager.loads)
    }

    @Test
    fun `falls back to default when no model satisfies the required capabilities`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model")),
            defaultModelId = "default-model",
        )

        val provider = router.route(stage, setOf(ModelCapability.Reasoning), modelId = null)

        assertEquals("managed:default-model", provider.id.value)
    }

    @Test
    fun `swap pins a model that then beats per-stage selection`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)),
            defaultModelId = "default-model",
        )

        router.swap("coder")
        assertEquals("coder", router.pinnedModel())

        // A stage explicitly asking for a different model is overridden by the pin.
        val provider = router.route(stage, emptySet(), modelId = "default-model")
        assertEquals("managed:coder", provider.id.value)
    }

    @Test
    fun `clearPin restores per-stage selection`(): Unit = runBlocking {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)),
            defaultModelId = "default-model",
        )

        router.swap("coder")
        router.clearPin()
        assertNull(router.pinnedModel())

        val provider = router.route(stage, emptySet(), modelId = "default-model")
        assertEquals("managed:default-model", provider.id.value)
    }

    @Test
    fun `swap to an unknown model throws and leaves the pin unset`() {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model")),
            defaultModelId = "default-model",
        )

        assertThrows<NoEligibleProviderException> { runBlocking { router.swap("ghost") } }
        assertNull(router.pinnedModel())
    }

    @Test
    fun `unknown explicit modelId throws NoEligibleProviderException`() {
        val manager = FakeModelManager()
        val router = ManagedInferenceRouter(
            manager = manager,
            descriptors = listOf(desc("default-model")),
            defaultModelId = "default-model",
        )

        assertThrows<NoEligibleProviderException> {
            runBlocking { router.route(stage, emptySet(), modelId = "ghost") }
        }
    }
}
