correx/testing/kernel/src/test/kotlin/ReplayInferenceProviderTest.kt

import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.FinishReason
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.TokenUsage
import com.correx.core.kernel.replay.ReplayInferenceProvider
import com.correx.core.kernel.retry.exception.ReplayArtifactMissingException
import com.correx.infrastructure.persistence.InMemoryEventStore
import com.correx.testing.fixtures.InferenceFixtures
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Assertions
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows

class ReplayInferenceProviderTest {

    private lateinit var eventStore: InMemoryEventStore
    private lateinit var replayProvider: ReplayInferenceProvider

    @BeforeEach
    fun setup() {
        eventStore = InMemoryEventStore()
        replayProvider = ReplayInferenceProvider(eventStore)
    }

    @Test
    fun `response carries correct tokensUsed and latencyMs when matching stageId`() = runTest {
        val sessionId = SessionId("s1")
        val stageId = StageId("stage1")
        val requestId = InferenceRequestId("req1")
        val providerId = ProviderId("test-provider")
        val tokensUsed = TokenUsage(promptTokens = 100, completionTokens = 200)
        val latencyMs = 500L

        // Record an inference completed event
        eventStore.append(
            InferenceFixtures.newInferenceCompletedEvent(
                requestId = requestId,
                sessionId = sessionId,
                stageId = stageId,
                providerId = providerId,
                tokensUsed = tokensUsed,
                latencyMs = latencyMs,
            ),
        )

        val request = InferenceFixtures.inferenceRequest(
            sessionId = sessionId,
            stageId = stageId,
            requestId = requestId,
        )

        val response = replayProvider.infer(request)

        Assertions.assertEquals(request.requestId, response.requestId)
        Assertions.assertEquals(tokensUsed, response.tokensUsed)
        Assertions.assertEquals(latencyMs, response.latencyMs)
        Assertions.assertEquals("", response.text)
        Assertions.assertEquals(FinishReason.Stop, response.finishReason)
    }

    @Test
    fun `ReplayArtifactMissingException thrown when event for different stageId`(): Unit = runBlocking {
        val sessionId = SessionId("s1")
        val stageId1 = StageId("stage1")
        val stageId2 = StageId("stage2")
        val requestId = InferenceRequestId("req1")
        val providerId = ProviderId("test-provider")

        // Record event for stage1
        eventStore.append(
            InferenceFixtures.newInferenceCompletedEvent(
                requestId = requestId,
                sessionId = sessionId,
                stageId = stageId1,
                providerId = providerId,
                tokensUsed = TokenUsage(promptTokens = 100, completionTokens = 200),
                latencyMs = 500L,
            ),
        )

        val request = InferenceFixtures.inferenceRequest(
            sessionId = sessionId,
            stageId = stageId2, // Different stage
            requestId = requestId,
        )

        assertThrows<ReplayArtifactMissingException> {
            replayProvider.infer(request)
        }
    }

    @Test
    fun `ReplayArtifactMissingException thrown when event store is empty`(): Unit = runBlocking {
        val sessionId = SessionId("s1")
        val stageId = StageId("stage1")
        val requestId = InferenceRequestId("req1")

        val request = InferenceFixtures.inferenceRequest(
            sessionId = sessionId,
            stageId = stageId,
            requestId = requestId,
        )

        assertThrows<ReplayArtifactMissingException> {
            replayProvider.infer(request)
        }
    }

    @Test
    fun `text is empty string by design`() = runTest {
        val sessionId = SessionId("s1")
        val stageId = StageId("stage1")
        val requestId = InferenceRequestId("req1")
        val providerId = ProviderId("test-provider")

        eventStore.append(
            InferenceFixtures.newInferenceCompletedEvent(
                requestId = requestId,
                sessionId = sessionId,
                stageId = stageId,
                providerId = providerId,
                tokensUsed = TokenUsage(promptTokens = 100, completionTokens = 50),
                latencyMs = 300L,
            ),
        )

        val request = InferenceFixtures.inferenceRequest(
            sessionId = sessionId,
            stageId = stageId,
            requestId = requestId,
        )

        val response = replayProvider.infer(request)

        Assertions.assertEquals("", response.text)
    }

    @Test
    fun `capabilities returns empty set`() {
        val capabilities = replayProvider.capabilities()

        Assertions.assertTrue(capabilities.isEmpty())
    }

    @Test
    fun `healthCheck returns Healthy`() = runTest {
        val health = replayProvider.healthCheck()

        Assertions.assertEquals(ProviderHealth.Healthy, health)
    }

    @Test
    fun `id is ReplayProvider`() {
        Assertions.assertEquals(ProviderId("replay-provider"), replayProvider.id)
    }
}