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

import com.correx.core.events.types.ProviderId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.FinishReason
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.PromptRenderer
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.ResponseFormat
import com.correx.core.inference.TokenUsage
import com.correx.core.inference.Tokenizer
import com.correx.infrastructure.inference.commons.ModelDescriptor
import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.engine.cio.CIO
import io.ktor.client.plugins.HttpTimeout
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.client.request.accept
import io.ktor.client.request.get
import io.ktor.client.request.post
import io.ktor.client.request.setBody
import io.ktor.client.statement.bodyAsText
import io.ktor.http.ContentType
import io.ktor.http.contentType
import io.ktor.serialization.kotlinx.json.json
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import org.slf4j.LoggerFactory

private const val DEFAULT_REQUEST_TIMEOUT_MS = 60_000L

private fun defaultHttpClient(): HttpClient = HttpClient(CIO) {
    install(ContentNegotiation) {
        json(
            Json {
                ignoreUnknownKeys = true
                isLenient = true
                encodeDefaults = false
            },
        )
    }
    install(HttpTimeout) {
        requestTimeoutMillis = DEFAULT_REQUEST_TIMEOUT_MS
    }
    install(ContentNegotiation) {
        json(
            Json {
                ignoreUnknownKeys = true
                explicitNulls = false
                encodeDefaults = true
            },
        )
    }
}
private val json = Json {
    ignoreUnknownKeys = true
    explicitNulls = false
    encodeDefaults = true
}

private val log = LoggerFactory.getLogger(LlamaCppInferenceProvider::class.java)

@Suppress("TooGenericExceptionCaught", "MagicNumber")
class LlamaCppInferenceProvider(
    private val descriptor: ModelDescriptor,
    private val baseUrl: String = "http://127.0.0.1:10000",
    private val httpClient: HttpClient = defaultHttpClient(),
) : InferenceProvider {

    private val modelId: String = descriptor.modelId

    override val id: ProviderId = ProviderId("llama-cpp:$modelId")
    override val name: String = "LlamaCpp ($modelId)"
    override val tokenizer: Tokenizer = LlamaCppTokenizer(baseUrl, httpClient)

    override suspend fun infer(request: InferenceRequest): InferenceResponse {
        val startTime = System.currentTimeMillis()

        val baseMessages = PromptRenderer.render(request.contextPack)
        val messages = baseMessages.map { msg ->
            LlamaCppChatMessage(
                role = msg.role,
                content = msg.content,
            )
        }

        val grammar = when (val fmt = request.responseFormat) {
            is ResponseFormat.Json -> GbnfGrammarConverter.convert(fmt.schema)
            ResponseFormat.Text -> null
        }

        val body = ChatCompletionRequest(
            model = modelId,
            messages = messages,
            temperature = request.generationConfig.temperature,
            topP = request.generationConfig.topP,
            maxTokens = request.generationConfig.maxTokens,
            stopSequences = request.generationConfig.stopSequences,
            seed = request.generationConfig.seed,
            stream = false,
            grammar = grammar,
            tools = request.tools.takeIf { it.isNotEmpty() },
        )

        val encoded = json.encodeToString(body)
        log.debug("sending request to llm: {}", encoded)

        val httpResponse = httpClient.post("$baseUrl/v1/chat/completions") {
            contentType(ContentType.Application.Json)
            accept(ContentType.Application.Json)
            setBody(encoded)
        }
        if (httpResponse.status.value !in 200..299) {
            val errorBody = httpResponse.bodyAsText()
            error("llama-server returned ${httpResponse.status.value} ${httpResponse.status.description}: $errorBody")
        }
        val response = httpResponse.body<ChatCompletionResponse>()

        log.debug("got response from llm: {}", response)

        val message = response.choices.first().message
        val finishReason = when (response.choices.first().finishReason.lowercase()) {
            "length" -> FinishReason.Length
            "tool_calls" -> FinishReason.ToolCall
            else -> FinishReason.Stop
        }

        return InferenceResponse(
            requestId = request.requestId,
            text = message.content ?: "",
            finishReason = finishReason,
            tokensUsed = TokenUsage(
                promptTokens = response.usage.promptTokens,
                completionTokens = response.usage.completionTokens,
            ),
            latencyMs = System.currentTimeMillis() - startTime,
            toolCalls = message.toolCalls,
        )
    }

    override suspend fun healthCheck(): ProviderHealth = try {
        val response = httpClient.get("$baseUrl/health")
        if (response.status.value in 200..299) {
            ProviderHealth.Healthy
        } else {
            ProviderHealth.Unavailable("Health check returned non-2xx status: ${response.status}")
        }
    } catch (e: Exception) {
        ProviderHealth.Unavailable("Health check failed: ${e.message}")
    }

    override fun capabilities(): Set<CapabilityScore> = descriptor.capabilities
}
