package com.correx.apps.server.inference

import com.correx.core.context.model.CompressionMetadata
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.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ResponseFormat
import java.util.UUID

suspend fun summarizeWithInference(
    prompt: String,
    sessionId: SessionId,
    inferenceRouter: InferenceRouter,
): String {
    val stageId = StageId("summarizer")
    val provider = inferenceRouter.route(stageId, setOf(ModelCapability.General))
    val contextPack = ContextPack(
        id = ContextPackId(UUID.randomUUID().toString()),
        sessionId = sessionId,
        stageId = stageId,
        layers = mapOf(
            ContextLayer.L0 to listOf(
                ContextEntry(
                    id = ContextEntryId(UUID.randomUUID().toString()),
                    layer = ContextLayer.L0,
                    content = prompt,
                    sourceType = "summarizer",
                    sourceId = "summarizer",
                    tokenEstimate = prompt.length / 4,
                    role = EntryRole.USER,
                )
            )
        ),
        budgetUsed = prompt.length / 4,
        budgetLimit = 512,
        compressionMetadata = CompressionMetadata(),
    )
    val request = InferenceRequest(
        requestId = InferenceRequestId(UUID.randomUUID().toString()),
        sessionId = sessionId,
        stageId = stageId,
        contextPack = contextPack,
        generationConfig = GenerationConfig(temperature = 0.3, topP = 0.9, maxTokens = 512),
        responseFormat = ResponseFormat.Text,
    )
    return provider.infer(request).text
}
