package com.correx.infrastructure.inference.openai

import com.correx.core.context.model.ContextPack
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.FinishReason
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.ModelCapability
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.client.request.HttpRequestData
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.headersOf
import io.ktor.serialization.kotlinx.json.json
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.Json
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertTrue

class OpenAiCompatInferenceProviderTest {

    private fun request() = InferenceRequest(
        requestId = InferenceRequestId("req1"),
        sessionId = SessionId("s1"),
        stageId = StageId("stage1"),
        contextPack = ContextPack(
            id = ContextPackId("cp1"),
            sessionId = SessionId("s1"),
            stageId = StageId("stage1"),
            layers = emptyMap(),
            budgetUsed = 0,
            budgetLimit = 4096,
        ),
        generationConfig = GenerationConfig(temperature = 0.2, topP = 0.9, maxTokens = 256),
    )

    private fun clientReturning(body: String, capture: (HttpRequestData) -> Unit = {}): HttpClient {
        val engine = MockEngine { req ->
            capture(req)
            respond(
                content = body,
                headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
            )
        }
        return HttpClient(engine) {
            install(ContentNegotiation) { json(Json { ignoreUnknownKeys = true }) }
        }
    }

    private fun provider(client: HttpClient) = OpenAiCompatInferenceProvider(
        modelId = "deepseek-ai/deepseek-v4-flash",
        baseUrl = "https://integrate.api.nvidia.com/v1",
        apiKey = "nvapi-test",
        capabilities = setOf(CapabilityScore(ModelCapability.General, 1.0)),
        idPrefix = "nim",
        httpClient = client,
    )

    @Test
    fun `infer sends bearer auth to chat completions and returns text`(): Unit = runBlocking {
        var seenAuth: String? = null
        var seenUrl: String? = null
        val client = clientReturning(
            """
            {"id":"x","choices":[{"message":{"role":"assistant","content":"hello there"},
            "finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":3,"total_tokens":14}}
            """.trimIndent(),
        ) { req ->
            seenAuth = req.headers[HttpHeaders.Authorization]
            seenUrl = req.url.toString()
        }

        val resp = provider(client).infer(request())

        assertEquals("Bearer nvapi-test", seenAuth)
        assertEquals("https://integrate.api.nvidia.com/v1/chat/completions", seenUrl)
        assertEquals("hello there", resp.text)
        assertEquals(FinishReason.Stop, resp.finishReason)
        assertEquals(11, resp.tokensUsed.promptTokens)
        assertEquals(3, resp.tokensUsed.completionTokens)
    }

    @Test
    fun `infer surfaces native tool calls as ToolCall finish`(): Unit = runBlocking {
        val client = clientReturning(
            """
            {"id":"x","choices":[{"message":{"role":"assistant","content":null,
            "tool_calls":[{"id":"c1","type":"function","function":{"name":"file_read","arguments":"{}"}}]},
            "finish_reason":"tool_calls"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}
            """.trimIndent(),
        )

        val resp = provider(client).infer(request())

        assertEquals(FinishReason.ToolCall, resp.finishReason)
        assertEquals(1, resp.toolCalls.size)
        assertEquals("file_read", resp.toolCalls.first().function.name)
    }

    @Test
    fun `infer salvages a tool call emitted as a JSON blob in content`(): Unit = runBlocking {
        val client = clientReturning(
            """
            {"id":"x","choices":[{"message":{"role":"assistant",
            "content":"{\"function\":{\"name\":\"file_read\",\"arguments\":\"{}\"}}"},
            "finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":9,"total_tokens":13}}
            """.trimIndent(),
        )

        val resp = provider(client).infer(request())

        assertEquals(FinishReason.ToolCall, resp.finishReason)
        assertEquals("file_read", resp.toolCalls.first().function.name)
        assertEquals("", resp.text)
    }

    @Test
    fun `salvageToolCalls leaves plain artifact JSON alone`() {
        assertTrue(salvageToolCalls("""{"title":"a plan","steps":[]}""").isEmpty())
        assertTrue(salvageToolCalls("not json").isEmpty())
    }
}
