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

import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.headersOf
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows

class LlamaCppEmbedderTest {

    @Test
    fun `embed returns float array with correct dimension from OpenAI format`(): Unit = runBlocking {
        val mockEngine = MockEngine { request ->
            val responseBody = """
                {
                    "data": [
                        {
                            "embedding": [0.1, 0.2, 0.3],
                            "index": 0
                        }
                    ]
                }
            """.trimIndent()
            respond(
                content = responseBody,
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString())
            )
        }

        val embeddingUrl = "http://localhost:11000"
        val httpClient = HttpClient(mockEngine)
        val embedder = LlamaCppEmbedder(
            url = embeddingUrl,
            modelId = "test-model",
            dimension = 3,
            httpClient = httpClient,
        )

        val result = embedder.embed("test text")

        assert(result.size == 3)
        assert(kotlin.math.abs(result[0] - 0.1f) < 0.0001f)
        assert(kotlin.math.abs(result[1] - 0.2f) < 0.0001f)
        assert(kotlin.math.abs(result[2] - 0.3f) < 0.0001f)
    }

    @Test
    fun `embed returns float array from simple embedding response format`(): Unit = runBlocking {
        val mockEngine = MockEngine { request ->
            val responseBody = """
                {
                    "embedding": [0.5, 0.6, 0.7]
                }
            """.trimIndent()
            respond(
                content = responseBody,
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString())
            )
        }

        val embeddingUrl = "http://localhost:11000"
        val httpClient = HttpClient(mockEngine)
        val embedder = LlamaCppEmbedder(
            url = embeddingUrl,
            modelId = null,
            dimension = 3,
            httpClient = httpClient,
        )

        val result = embedder.embed("test text")

        assert(result.size == 3)
        assert(kotlin.math.abs(result[0] - 0.5f) < 0.0001f)
        assert(kotlin.math.abs(result[1] - 0.6f) < 0.0001f)
        assert(kotlin.math.abs(result[2] - 0.7f) < 0.0001f)
    }

    @Test
    fun `embed parses llama-cpp nested pooling-mean array format`(): Unit = runBlocking {
        // llama.cpp /embedding with --pooling mean: [{"index":0,"embedding":[[...]]}]
        val mockEngine = MockEngine {
            respond(
                content = """[{"index":0,"embedding":[[0.5,0.6,0.7]]}]""",
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString())
            )
        }
        val embedder = LlamaCppEmbedder(
            url = "http://localhost:11000",
            dimension = 3,
            httpClient = HttpClient(mockEngine),
        )

        val result = embedder.embed("test text")

        assert(result.size == 3)
        assert(kotlin.math.abs(result[0] - 0.5f) < 0.0001f)
        assert(kotlin.math.abs(result[2] - 0.7f) < 0.0001f)
    }

    @Test
    fun `embed throws on dimension mismatch`(): Unit = runBlocking {
        val mockEngine = MockEngine { request ->
            val responseBody = """
                {
                    "embedding": [0.1, 0.2, 0.3, 0.4]
                }
            """.trimIndent()
            respond(
                content = responseBody,
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString())
            )
        }

        val embeddingUrl = "http://localhost:11000"
        val httpClient = HttpClient(mockEngine)
        val embedder = LlamaCppEmbedder(
            url = embeddingUrl,
            dimension = 3,
            httpClient = httpClient,
        )

        assertThrows<IllegalStateException> {
            runBlocking {
                embedder.embed("test text")
            }
        }
    }

    @Test
    fun `embed throws on unexpected response format`(): Unit = runBlocking {
        val mockEngine = MockEngine { request ->
            val responseBody = """
                {
                    "unexpected_field": "value"
                }
            """.trimIndent()
            respond(
                content = responseBody,
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString())
            )
        }

        val embeddingUrl = "http://localhost:11000"
        val httpClient = HttpClient(mockEngine)
        val embedder = LlamaCppEmbedder(
            url = embeddingUrl,
            dimension = 3,
            httpClient = httpClient,
        )

        assertThrows<IllegalStateException> {
            runBlocking {
                embedder.embed("test text")
            }
        }
    }
}
