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

import com.correx.core.inference.Embedder
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.post
import io.ktor.client.request.setBody
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 const val ERROR_BODY_PREVIEW_LENGTH = 200

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

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

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

class LlamaCppEmbedder(
    private val url: String,
    private val modelId: String? = null,
    override val dimension: Int = 1536,
    private val httpClient: HttpClient = defaultHttpClient(),
) : Embedder {

    override suspend fun embed(text: String): FloatArray {
        val request = EmbeddingRequest(
            content = text,
            model = modelId,
        )

        val encoded = json.encodeToString(request)
        log.debug("sending embedding request to llama.cpp: {}", encoded)

        val responseBody = try {
            httpClient.post("$url/embedding") {
                contentType(ContentType.Application.Json)
                accept(ContentType.Application.Json)
                setBody(encoded)
            }.body<String>()
        } catch (e: Exception) {
            throw IllegalStateException("Failed to call embedding endpoint at $url/embedding: ${e.message}", e)
        }

        log.debug("got embedding response from llama.cpp: {}", responseBody)

        val embedding = parseEmbeddingResponse(responseBody)

        if (embedding.size != dimension) {
            throw IllegalStateException(
                "Embedding dimension mismatch: expected $dimension but got ${embedding.size}"
            )
        }

        return embedding.toFloatArray()
    }

    private fun parseEmbeddingResponse(responseBody: String): List<Float> {
        // Try OpenAI-compatible format first: {"data": [{"embedding": [...], "index": 0}]}
        val openAiResult = runCatching {
            val openAiResponse = json.decodeFromString<OpenAiEmbeddingResponse>(responseBody)
            if (openAiResponse.data.isNotEmpty()) {
                openAiResponse.data.first().embedding
            } else {
                emptyList<Float>()
            }
        }.getOrNull()
        if (openAiResult != null && openAiResult.isNotEmpty()) {
            return openAiResult
        }

        // Try simple format: {"embedding": [...]}
        val simpleResult = runCatching {
            val simpleResponse = json.decodeFromString<EmbeddingResponse>(responseBody)
            simpleResponse.embedding ?: emptyList<Float>()
        }.getOrNull()
        if (simpleResult != null && simpleResult.isNotEmpty()) {
            return simpleResult
        }

        // Try array format: [{"embedding": [...]}]
        val arrayResult = runCatching {
            val arrayResponse = json.decodeFromString<List<EmbeddingArrayResponse>>(responseBody)
            if (arrayResponse.isNotEmpty()) {
                arrayResponse.first().embedding
            } else {
                emptyList<Float>()
            }
        }.getOrNull()
        if (arrayResult != null && arrayResult.isNotEmpty()) {
            return arrayResult
        }

        // llama.cpp --pooling mean: [{"embedding": [[...]]}] (nested list-of-one-vector).
        val nestedResult = runCatching {
            json.decodeFromString<List<EmbeddingNestedArrayResponse>>(responseBody)
                .firstOrNull()?.embedding?.firstOrNull().orEmpty()
        }.getOrNull()
        if (nestedResult != null && nestedResult.isNotEmpty()) {
            return nestedResult
        }

        val preview = responseBody.take(ERROR_BODY_PREVIEW_LENGTH)
        throw IllegalStateException(
            "Unexpected embedding response shape from $url. " +
                "Response body (first $ERROR_BODY_PREVIEW_LENGTH chars): $preview"
        )
    }
}
