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.core.inference.ToolCallFunction
import com.correx.core.inference.ToolCallRequest
import com.correx.core.inference.ToolDefinition
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

// Aligned with the orchestrator's per-stage timeout (1_800_000ms). A slow local reasoning model
// (large prompt + a multi-thousand-token reasoning budget) can take >10min for a single call;
// a shorter HTTP timeout here fails the call before the stage timeout ever applies.
private const val DEFAULT_REQUEST_TIMEOUT_MS = 1_800_000L

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

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

// salvageToolCalls parses a tool call the model wrote into `content` as raw JSON (single object
// or array) instead of the native tool_calls array. Real artifact JSON has no `function` field,
// so it fails to decode and is left alone — only genuine tool-call shapes are recovered.
internal fun salvageToolCalls(content: String?): List<ToolCallRequest> {
    val s = content?.trim().orEmpty()
    if (s.isEmpty() || (s.first() != '{' && s.first() != '[')) return emptyList()
    val strict = runCatching { json.decodeFromString<List<ToolCallRequest>>(s) }
        .recoverCatching { listOf(json.decodeFromString<ToolCallRequest>(s)) }
        .getOrDefault(emptyList())
        .filter { it.function.name.isNotBlank() }
    return strict.ifEmpty { salvageLenient(s) }
}

private val NAME_RE = Regex("\"name\"\\s*:\\s*\"([^\"]+)\"")

// Small local models often emit a tool call into `content` but truncate it — an unterminated
// arguments string and missing closing braces — so strict JSON decode fails. Only the `function`
// shape (name + arguments) matters to the orchestrator; recover those by hand. Gated on a tool-call
// marker so genuine artifact JSON (no `function`/`arguments` field) is never misread as a call.
private fun salvageLenient(s: String): List<ToolCallRequest> {
    if (!s.contains("\"function\"") && !s.contains("\"arguments\"")) return emptyList()
    val name = NAME_RE.find(s)?.groupValues?.get(1) ?: return emptyList()
    val argsIdx = s.indexOf("\"arguments\"")
    val args = if (argsIdx >= 0) {
        // The arguments value is a JSON string holding escaped JSON; unescape, then take the
        // first balanced object so trailing truncation garbage is dropped.
        val unescaped = s.substring(argsIdx).replace("\\\"", "\"").replace("\\\\", "\\")
        firstBalancedObject(unescaped) ?: "{}"
    } else {
        "{}"
    }
    return listOf(ToolCallRequest(function = ToolCallFunction(name = name, arguments = args)))
}

// Returns the first brace-balanced JSON object substring, or null if none closes (true truncation).
private fun firstBalancedObject(s: String): String? {
    val start = s.indexOf('{')
    if (start < 0) return null
    var depth = 0
    var inStr = false
    var esc = false
    for (i in start until s.length) {
        val c = s[i]
        when {
            esc -> esc = false
            c == '\\' -> esc = true
            c == '"' -> inStr = !inStr
            !inStr && c == '{' -> depth++
            !inStr && c == '}' -> {
                depth--
                if (depth == 0) return s.substring(start, i + 1)
            }
        }
    }
    return null
}

@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 ->
            val toolCalls = if (msg.role == "assistant") {
                parseAssistantToolCalls(msg.content)
            } else emptyList()
            LlamaCppChatMessage(
                role = msg.role,
                content = if (toolCalls.isNotEmpty()) null else msg.content,
                reasoningContent = msg.reasoning,
                toolCalls = toolCalls,
                toolCallId = msg.toolCallId,
            )
        }

        val tools = request.tools.takeIf { it.isNotEmpty() }

        // #291: clamp the stage completion cap to the window's real runway (see clampMaxTokens).
        val effectiveMaxTokens = clampMaxTokens(request.generationConfig.maxTokens, messages, tools)

        // llama.cpp rejects requests that carry BOTH a custom grammar and tools
        // ("Cannot use custom grammar constraints with tools"). When a stage has tools,
        // we drop the grammar and rely on post-hoc schema validation + retry to keep the
        // artifact constrained (invariant #7: LLM output is untrusted until validated).
        // TODO: replace this implicit "free-form JSON, validate after" path with a
        //  first-class artifact-emission tool — the LLM calls e.g. emit_artifact(path,
        //  content, description) the same way it calls file_read/file_write, so artifact
        //  production is an explicit, schema-checked tool call rather than a grammar
        //  constraint that can't coexist with other tools.
        val grammar = when (val fmt = request.responseFormat) {
            is ResponseFormat.Json -> if (tools == null) GbnfGrammarConverter.convert(fmt.schema) else null
            ResponseFormat.Text -> null
        }

        val body = ChatCompletionRequest(
            model = modelId,
            messages = messages,
            temperature = request.generationConfig.temperature,
            topP = request.generationConfig.topP,
            maxTokens = effectiveMaxTokens,
            stopSequences = request.generationConfig.stopSequences,
            seed = request.generationConfig.seed,
            topK = request.generationConfig.topK,
            minP = request.generationConfig.minP,
            repeatPenalty = request.generationConfig.repeatPenalty,
            stream = false,
            grammar = grammar,
            tools = tools,
        )

        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 choice = response.choices.firstOrNull()
            ?: error("llama-server returned no choices in the completion response")
        val message = choice.message
        // Some local models emit the tool call as a JSON blob in `content` instead of the native
        // tool_calls array; salvage it so the orchestrator gets a real call rather than treating
        // the blob as a (failing) artifact and retry-looping to exhaustion.
        val salvaged = if (message.toolCalls.isEmpty()) salvageToolCalls(message.content) else emptyList()
        val toolCalls = message.toolCalls.ifEmpty { salvaged }
        val finishReason = when {
            toolCalls.isNotEmpty() -> FinishReason.ToolCall
            choice.finishReason.lowercase() == "length" -> FinishReason.Length
            else -> FinishReason.Stop
        }

        return InferenceResponse(
            requestId = request.requestId,
            text = if (salvaged.isNotEmpty()) "" else message.content ?: "",
            finishReason = finishReason,
            tokensUsed = TokenUsage(
                promptTokens = response.usage.promptTokens,
                completionTokens = response.usage.completionTokens,
            ),
            latencyMs = System.currentTimeMillis() - startTime,
            toolCalls = toolCalls,
            reasoning = message.reasoningContent?.takeIf { it.isNotBlank() },
        )
    }

    // Per-message chat-template framing (role tags, delimiters) the client can't see but the server
    // adds; a small fixed estimate is enough headroom for a safety clamp. ponytail: constant, not a
    // per-family table — bump if a template proves heavier.
    private val templateTokensPerMessage = 8
    private val safetyReserveTokens = 512
    private val minCompletionTokens = 256

    /**
     * #291: returns min(requestedCap, contextSize − promptTokens − toolSchema − templateOverhead −
     * reserve), floored so a nearly-full window still asks for *something* rather than a negative or
     * absurd allowance. Prompt/tool tokens are counted with the model's own tokenizer (one extra
     * /tokenize round-trip); a char/4 estimate is the fallback if that call fails. Live-path only —
     * deterministic replay never calls [infer], so nothing here needs to be an event.
     */
    private suspend fun clampMaxTokens(
        requestedCap: Int,
        messages: List<LlamaCppChatMessage>,
        tools: List<ToolDefinition>?,
    ): Int {
        val countable = buildString {
            messages.forEach { m ->
                append(m.role).append('\n')
                m.content?.let { append(it).append('\n') }
                m.reasoningContent?.let { append(it).append('\n') }
                m.toolCalls.forEach { append(it.function.name).append(it.function.arguments).append('\n') }
            }
            tools?.let { append(json.encodeToString(it)) }
        }
        val promptTokens = runCatching { tokenizer.countTokens(countable) }
            .getOrElse { countable.length / 4 }
        val overhead = messages.size * templateTokensPerMessage
        val runway = descriptor.contextSize - promptTokens - overhead - safetyReserveTokens
        if (runway < minCompletionTokens) {
            log.warn(
                "context runway low: ctx={} prompt~{} overhead={} reserve={} runway={} (cap was {})",
                descriptor.contextSize, promptTokens, overhead, safetyReserveTokens, runway, requestedCap,
            )
        }
        return requestedCap.coerceAtMost(runway).coerceAtLeast(0)
    }

    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}")
    }

    /**
     * Parses an assistant message's content back into [ToolCallRequest] objects.
     *
     * The context builder's FORMAT_COMPRESS stage converts raw ToolCallRequest JSON into a compact
     * dotted-line format (id = ..., function.name = ...). If compression is disabled, the content
     * is still valid JSON. This handles both forms so the Gemma4 Jinja template sees the native
     * `tool_calls` array — which it requires to render `reasoning_content` as a `<|channel>thought`
     * block (without `tool_calls` the reasoning field is silently dropped and the model re-enters
     * thought mode from scratch, causing the degenerate `<|channel>thought`~16384-token loop).
     */
    private fun parseAssistantToolCalls(content: String): List<ToolCallRequest> {
        // Try JSON parse first (FORMAT_COMPRESS disabled — raw ToolCallRequest JSON).
        if (content.startsWith("{\"id\"")) {
            runCatching {
                return listOf(json.decodeFromString<ToolCallRequest>(content))
            }
        }
        // Try dotted-line format (FORMAT_COMPRESS enabled).
        if (!content.contains("function.name")) return emptyList()
        val lines = content.split('\n')
        var id: String? = null
        var name: String? = null
        var arguments: String? = null
        for (line in lines) {
            when {
                line.startsWith("id = ") -> id = line.removePrefix("id = ").trim()
                line.startsWith("function.name = ") -> name = line.removePrefix("function.name = ").trim()
                line.startsWith("function.arguments = ") ->
                    arguments = line.removePrefix("function.arguments = ").trim()
            }
        }
        return if (name != null && arguments != null) {
            listOf(ToolCallRequest(id = id, function = ToolCallFunction(name = name, arguments = arguments)))
        } else emptyList()
    }

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