package com.correx.testing.contracts.model.inference

import com.correx.core.context.model.CompressionMetadata
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.GenerationConfig
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.testing.fixtures.inference.MockInferenceProvider
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class InferenceProviderContractTest {

    fun provider(): InferenceProvider = MockInferenceProvider()

    private fun request() = InferenceRequest(
        requestId = InferenceRequestId("req-1"),
        sessionId = SessionId("session-1"),
        stageId = StageId("stage-1"),
        contextPack = emptyContextPack(),           // test helper — returns minimal stub
        generationConfig = GenerationConfig(
            temperature = 0.0,
            topP = 1.0,
            maxTokens = 128,
        ),
    )

    // ── Successful inference ──────────────────────────────────────────────────

    @Test
    fun `infer returns response with matching requestId`(): Unit = runBlocking {
        val p = provider()
        val req = request()
        val resp = p.infer(req)
        assertEquals(req.requestId, resp.requestId)
    }

    @Test
    fun `infer returns non-blank text`(): Unit = runBlocking {
        val resp = provider().infer(request())
        assertTrue(resp.text.isNotBlank())
    }

    @Test
    fun `infer tracks token usage`(): Unit = runBlocking {
        val resp = provider().infer(request())
        assertTrue(resp.tokensUsed.totalTokens > 0)
    }

    @Test
    fun `infer latency is non-negative`(): Unit = runBlocking {
        val resp = provider().infer(request())
        assertTrue(resp.latencyMs >= 0)
    }

    // ── Health ────────────────────────────────────────────────────────────────

    @Test
    fun `healthCheck returns non-null`(): Unit = runBlocking {
        assertNotNull(provider().healthCheck())
    }

    // ── Capabilities ──────────────────────────────────────────────────────────

    @Test
    fun `capabilities is non-empty`() {
        assertTrue(provider().capabilities().isNotEmpty())
    }

    @Test
    fun `all capability scores are in range 0 to 1`() {
        provider().capabilities().forEach { cs ->
            assertTrue(cs.score in 0.0..1.0, "score out of range: $cs")
        }
    }

    // ── Cancellation ──────────────────────────────────────────────────────────

    // A slow provider — required so cancellation actually fires mid-flight rather than
    // racing a 0ms infer that has already completed by the time cancel() is called.
    private fun slowProvider(): InferenceProvider = MockInferenceProvider(artificialDelayMs = 1_000L)

    @Test
    fun `infer respects coroutine cancellation`(): Unit = runBlocking {
        val p = slowProvider()
        val started = System.currentTimeMillis()
        val job = launch { p.infer(request()) }
        // give the coroutine a moment to enter delay()
        while (job.isActive && System.currentTimeMillis() - started < 50) { /* spin */ }
        job.cancel()
        job.join()
        assertFalse(job.isActive)
        // Must complete well before the 1000ms artificial delay would have elapsed.
        assertTrue(System.currentTimeMillis() - started < 500, "cancellation did not interrupt mid-flight")
    }

    @Test
    fun `infer cancellation completes within 500ms`(): Unit = runBlocking {
        val p = slowProvider()
        val started = System.currentTimeMillis()
        val job = launch { p.infer(request()) }
        job.cancel()
        withTimeout(500) { job.join() }
        assertTrue(System.currentTimeMillis() - started < 500, "join exceeded budget after cancel")
    }

    // ── Tokenizer ─────────────────────────────────────────────────────────────

    @Test
    fun `tokenizer countTokens is consistent with tokenize`(): Unit = runBlocking {
        val t = provider().tokenizer
        val text = "hello world"
        assertEquals(t.tokenize(text).size, t.countTokens(text))
    }

    private fun emptyContextPack(): ContextPack = ContextPack(
        id = ContextPackId("cp1"),
        sessionId = SessionId("session-1"),
        stageId = StageId("stage-1"),
        layers = emptyMap(),
        budgetUsed = 10,
        budgetLimit = 20,
        compressionMetadata = CompressionMetadata(),
    )
}
