package com.correx.infrastructure

import com.correx.core.config.ModelConfig
import com.correx.core.config.ModelsSettings
import com.correx.core.events.EventDispatcher
import com.correx.core.inference.ModelCapability
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.persistence.InMemoryEventStore
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.headersOf
import org.junit.jupiter.api.Test
import kotlin.test.assertEquals
import kotlin.test.assertNotNull

class InfrastructureModuleModelTest {

    @Test
    fun `modelConfigToDescriptor maps fields correctly`() {
        val config = ModelConfig(
            id = "my-model",
            modelPath = "/models/my-model.gguf",
            contextSize = 4096,
            capabilities = mapOf("General" to 1.0, "Coding" to 0.8),
        )

        val descriptor = InfrastructureModule.modelConfigToDescriptor(config)

        assertEquals("my-model", descriptor.modelId)
        assertEquals("/models/my-model.gguf", descriptor.modelPath)
        assertEquals(4096, descriptor.contextSize)
        assertEquals(ResidencyMode.PERSISTENT, descriptor.residencyMode)
        assertEquals(2, descriptor.capabilities.size)
        val generalScore = descriptor.capabilities.firstOrNull { it.capability == ModelCapability.General }
        assertNotNull(generalScore)
        assertEquals(1.0, generalScore!!.score)
        val codingScore = descriptor.capabilities.firstOrNull { it.capability == ModelCapability.Coding }
        assertNotNull(codingScore)
        assertEquals(0.8, codingScore!!.score)
    }

    @Test
    fun `modelConfigToDescriptor flattens params into extraArgs for llama-server`() {
        val config = ModelConfig(
            id = "mtp-model",
            modelPath = "/models/main.gguf",
            params = linkedMapOf(
                "-md" to "/models/draft.gguf",
                "-ngld" to "99",
                "--spec-type" to "draft-mtp",
                "--spec-draft-n-max" to "4",
            ),
        )

        val descriptor = InfrastructureModule.modelConfigToDescriptor(config)

        assertEquals(
            listOf("-md", "/models/draft.gguf", "-ngld", "99", "--spec-type", "draft-mtp", "--spec-draft-n-max", "4"),
            descriptor.extraArgs,
        )
    }

    @Test
    fun `modelConfigToDescriptor handles empty capabilities`() {
        val config = ModelConfig(
            id = "bare-model",
            modelPath = "/models/bare.gguf",
        )

        val descriptor = InfrastructureModule.modelConfigToDescriptor(config)

        assertEquals("bare-model", descriptor.modelId)
        assertEquals(24_576, descriptor.contextSize)
        assertEquals(0, descriptor.capabilities.size)
    }

    @Test
    fun `createModelManager returns DefaultModelManager with correct settings`() {
        val settings = ModelsSettings(
            defaultModel = "test-model",
            llamaServerBin = "/custom/llama-server",
            host = "10.0.0.1",
            port = 12345,
        )
        val eventStore = InMemoryEventStore()
        val dispatcher = EventDispatcher(eventStore)
        val mockEngine = MockEngine {
            respond(
                content = """{"status":"healthy"}""",
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
            )
        }
        val httpClient = HttpClient(mockEngine)

        val manager = InfrastructureModule.createModelManager(
            settings = settings,
            eventStore = eventStore,
            eventDispatcher = dispatcher,
            httpClient = httpClient,
        )

        assertNotNull(manager)
        // currentModel is null before any load
        assertEquals(null, manager.currentModel())
    }
}
