package com.correx.infrastructure.inference.openai

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.TokenUsage
import com.correx.core.inference.Tokenizer
import com.correx.core.inference.ToolCallRequest
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.header
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.HttpHeaders
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 = 600_000L

private val json = Json {
    ignoreUnknownKeys = true
    explicitNulls = false
    encodeDefaults = true
}

private fun defaultHttpClient(): HttpClient = HttpClient(CIO) {
    install(ContentNegotiation) { json(json) }
    install(HttpTimeout) { requestTimeoutMillis = DEFAULT_REQUEST_TIMEOUT_MS }
}

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

// Some models emit a tool call as a raw 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. Mirrors the llama.cpp provider's recovery path.
internal fun salvageToolCalls(content: String?): List<ToolCallRequest> {
    val s = content?.trim().orEmpty()
    if (s.isEmpty() || (s.first() != '{' && s.first() != '[')) return emptyList()
    return runCatching { json.decodeFromString<List<ToolCallRequest>>(s) }
        .recoverCatching { listOf(json.decodeFromString<ToolCallRequest>(s)) }
        .getOrDefault(emptyList())
        .filter { it.function.name.isNotBlank() }
}

/**
 * [InferenceProvider] for OpenAI-compatible chat APIs reached over HTTPS with a Bearer key —
 * NVIDIA NIM (`https://integrate.api.nvidia.com/v1`), OpenAI, vLLM, etc. Inference is fully
 * remote: no local process, no GPU, no GGUF. [baseUrl] must already include the API version
 * segment (e.g. `.../v1`); this class appends `/chat/completions` and `/models`.
 *
 * @param idPrefix label for the [ProviderId] (e.g. `nim`, `openai`).
 * @param apiKey Bearer token; sent as `Authorization: Bearer <key>` when non-blank.
 */
@Suppress("TooGenericExceptionCaught", "MagicNumber")
class OpenAiCompatInferenceProvider(
    private val modelId: String,
    private val baseUrl: String,
    private val apiKey: String,
    private val capabilities: Set<CapabilityScore>,
    private val idPrefix: String = "openai",
    private val httpClient: HttpClient = defaultHttpClient(),
) : InferenceProvider {

    override val id: ProviderId = ProviderId("$idPrefix:$modelId")
    override val name: String = "OpenAI-compatible ($idPrefix:$modelId)"
    override val tokenizer: Tokenizer = HeuristicTokenizer()

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

        val messages = PromptRenderer.render(request.contextPack).map { msg ->
            // Strict OpenAI-compatible servers reject a `tool` role that isn't bound to a prior
            // tool_calls turn (we don't carry tool_call_ids through the context pack). Fold any
            // non-standard role into a user turn so the model still sees the content.
            if (msg.role in STANDARD_ROLES) {
                OpenAiChatMessage(role = msg.role, content = msg.content, reasoningContent = msg.reasoning)
            } else {
                OpenAiChatMessage(role = "user", content = "[${msg.role}] ${msg.content}")
            }
        }

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

        val body = OpenAiChatCompletionRequest(
            model = modelId,
            messages = messages,
            temperature = request.generationConfig.temperature,
            topP = request.generationConfig.topP,
            maxTokens = request.generationConfig.maxTokens,
            stopSequences = request.generationConfig.stopSequences.ifEmpty { null },
            seed = request.generationConfig.seed,
            topK = request.generationConfig.topK,
            minP = request.generationConfig.minP,
            repeatPenalty = request.generationConfig.repeatPenalty,
            stream = false,
            tools = tools,
        )

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

        val httpResponse = httpClient.post("$baseUrl/chat/completions") {
            contentType(ContentType.Application.Json)
            accept(ContentType.Application.Json)
            if (apiKey.isNotBlank()) header(HttpHeaders.Authorization, "Bearer $apiKey")
            setBody(encoded)
        }
        if (httpResponse.status.value !in 200..299) {
            val errorBody = httpResponse.bodyAsText()
            error("$idPrefix returned ${httpResponse.status.value} ${httpResponse.status.description}: $errorBody")
        }
        val response = httpResponse.body<OpenAiChatCompletionResponse>()
        log.debug("got response from {}: {}", id.value, response)

        val choice = response.choices.firstOrNull()
            ?: error("$idPrefix returned no choices")
        val message = choice.message
        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 ?: 0,
                completionTokens = response.usage?.completionTokens ?: 0,
            ),
            latencyMs = System.currentTimeMillis() - startTime,
            toolCalls = toolCalls,
            reasoning = message.reasoningContent?.takeIf { it.isNotBlank() },
        )
    }

    override suspend fun healthCheck(): ProviderHealth = try {
        val response = httpClient.get("$baseUrl/models") {
            if (apiKey.isNotBlank()) header(HttpHeaders.Authorization, "Bearer $apiKey")
        }
        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> = capabilities

    private companion object {
        private val STANDARD_ROLES = setOf("system", "user", "assistant")
    }
}
