package com.correx.infrastructure.inference.llama.cpp

import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.EntryRole
import com.correx.core.events.types.ContextEntryId
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.InferenceRequest
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ResidencyMode
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.serialization.kotlinx.json.json
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.content.TextContent
import io.ktor.http.headersOf
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Test
import kotlin.test.assertEquals

// #291: the provider must clamp the requested completion cap against the model context window so a
// nearly-full prompt never asks for an impossible allowance (the traced Gemma4 truncation).
class LlamaCppMaxTokensClampTest {

    private val promptTokenCount = 1000
    private val reserve = 512
    private val overheadPerMessage = 8

    private fun request(cap: Int) = InferenceRequest(
        requestId = InferenceRequestId("req-1"),
        sessionId = SessionId("s-1"),
        stageId = StageId("stage-1"),
        contextPack = ContextPack(
            id = ContextPackId("pack-1"),
            sessionId = SessionId("s-1"),
            stageId = StageId("stage-1"),
            layers = mapOf(
                ContextLayer.L1 to listOf(
                    ContextEntry(
                        id = ContextEntryId("e-1"),
                        layer = ContextLayer.L1,
                        content = "do the thing",
                        sourceType = "stagePrompt",
                        sourceId = "e-1",
                        tokenEstimate = 4,
                        role = EntryRole.USER,
                    ),
                ),
            ),
            budgetUsed = 4,
            budgetLimit = 24_576,
        ),
        generationConfig = GenerationConfig(temperature = 1.0, topP = 1.0, maxTokens = cap),
    )

    private fun provider(contextSize: Int, capture: (Int) -> Unit): LlamaCppInferenceProvider {
        val engine = MockEngine { req ->
            when {
                req.url.encodedPath.endsWith("/tokenize") -> respond(
                    content = """{"tokens":[${(1..promptTokenCount).joinToString(",")}]}""",
                    headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
                )
                else -> {
                    val body = (req.body as TextContent).text
                    val maxTokens = Regex("\"max_tokens\":(\\d+)").find(body)!!.groupValues[1].toInt()
                    capture(maxTokens)
                    val usage =
                        """"usage":{"prompt_tokens":$promptTokenCount,"completion_tokens":1,"total_tokens":1}"""
                    respond(
                        content = """{"id":"x","choices":[{"message":{"role":"assistant",""" +
                            """"content":"ok"},"finish_reason":"stop"}],$usage}""",
                        headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
                    )
                }
            }
        }
        val descriptor = ModelDescriptor(
            modelId = "test",
            modelPath = "/dev/null",
            residencyMode = ResidencyMode.EPHEMERAL,
            contextSize = contextSize,
        )
        val client = HttpClient(engine) {
            install(ContentNegotiation) { json(kotlinx.serialization.json.Json { ignoreUnknownKeys = true }) }
        }
        return LlamaCppInferenceProvider(descriptor, "http://localhost:10000", client)
    }

    @Test
    fun `clamps completion cap to remaining context runway`(): Unit = runBlocking {
        var sent = -1
        // 1 message -> overhead 8. runway = 2000 - 1000 - 8 - 512 = 480, below the 24_576 cap.
        provider(contextSize = 2000) { sent = it }.infer(request(cap = 24_576))
        assertEquals(2000 - promptTokenCount - overheadPerMessage - reserve, sent)
    }

    @Test
    fun `keeps the stage cap when the window has ample runway`(): Unit = runBlocking {
        var sent = -1
        provider(contextSize = 100_000) { sent = it }.infer(request(cap = 24_576))
        assertEquals(24_576, sent)
    }
}
