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

import com.correx.core.inference.ChatMessage
import com.correx.core.inference.ToolCallRequest
import com.correx.core.inference.ToolDefinition
import kotlinx.serialization.EncodeDefault
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable

@Serializable
data class ChatCompletionRequest(
    val model: String,
    val messages: List<LlamaCppChatMessage>,
    val temperature: Double,
    @SerialName("top_p") val topP: Double,
    @SerialName("max_tokens") val maxTokens: Int,
    @SerialName("stop") val stopSequences: List<String> = emptyList(),
    val seed: Long? = null,
    // EncodeDefault.NEVER overrides the class-level encodeDefaults=true so an unset (null) sampling
    // knob is omitted from the JSON entirely, letting the model keep its own default (rather than
    // sending "top_k": null, which llama.cpp may reject or misread).
    @EncodeDefault(EncodeDefault.Mode.NEVER) @SerialName("top_k") val topK: Int? = null,
    @EncodeDefault(EncodeDefault.Mode.NEVER) @SerialName("min_p") val minP: Double? = null,
    @EncodeDefault(EncodeDefault.Mode.NEVER) @SerialName("repeat_penalty") val repeatPenalty: Double? = null,
    val stream: Boolean = false,
    val grammar: String? = null,
    val tools: List<ToolDefinition>? = null,
)

@Serializable
data class LlamaCppChatMessage(
    val role: String,
    val content: String?,
    @SerialName("reasoning_content") val reasoningContent: String? = null,
    @SerialName("tool_calls") val toolCalls: List<ToolCallRequest> = emptyList(),
    @SerialName("tool_call_id") val toolCallId: String? = null,
)

@Serializable
data class ChatCompletionResponse(
    val id: String,
    val choices: List<Choice>,
    val usage: Usage,
)

@Serializable
data class Choice(
    val message: LlamaCppChatMessage,
    @SerialName("finish_reason") val finishReason: String,
)

@Serializable
data class Usage(
    @SerialName("prompt_tokens") val promptTokens: Int,
    @SerialName("completion_tokens") val completionTokens: Int,
    @SerialName("total_tokens") val totalTokens: Int,
)

@Serializable
data class TokenizeRequest(val content: String)

@Serializable
data class TokenizeResponse(val tokens: List<Int>)
