import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceStartedEvent
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.InferenceProjector
import com.correx.core.inference.InferenceRecord
import com.correx.core.inference.InferenceState
import com.correx.core.inference.InferenceStatus
import com.correx.core.inference.TokenUsage
import com.correx.core.utils.TypeId
import com.correx.testing.fixtures.EventFixtures.stored
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test

class InferenceProjectorTest {

    private lateinit var projector: InferenceProjector

    @BeforeEach
    fun setup() {
        projector = InferenceProjector()
    }

    @Test
    fun `should initialize state with empty records`() {
        val initialState = projector.initial()
        assertEquals(0, initialState.records.size)
    }

    @Test
    fun `should start a new record on InferenceStartedEvent`() {
        val initialState = projector.initial()
        val event = stored(
            payload = InferenceStartedEvent(
                requestId = InferenceRequestId("req1"),
                sessionId = SessionId("sess1"),
                stageId = StageId("stageA"),
                providerId = ProviderId("providerX"),
                promptArtifactId = TypeId("00".repeat(32)),
            ),
        )
        val newState = projector.apply(initialState, event)

        assertEquals(1, newState.records.size)
        val record = newState.records.first()
        assertEquals(InferenceStatus.STARTED, record.status)
        assertEquals("req1", record.requestId.value)
    }

    @Test
    fun `should complete a record on InferenceCompletedEvent`() {
        val startedRecord = InferenceRecord(
            requestId = InferenceRequestId("req1"),
            sessionId = SessionId("sess1"),
            stageId = StageId("stageA"),
            providerId = ProviderId("providerX"),
            status = InferenceStatus.STARTED,
        )
        val initialState = InferenceState(listOf(startedRecord))

        val event = stored(
            payload = InferenceCompletedEvent(
                requestId = InferenceRequestId("req1"),
                sessionId = SessionId("sess1"),
                stageId = StageId("stageA"),
                providerId = ProviderId("providerX"),
                tokensUsed = TokenUsage(50, 50),
                latencyMs = 500L,
                responseArtifactId = TypeId("00".repeat(32)),
            ),
        )
        val newState = projector.apply(initialState, event)

        assertEquals(1, newState.records.size)
        val record = newState.records.first()
        assertEquals(InferenceStatus.COMPLETED, record.status)
        assertEquals(100, record.tokensUsed?.totalTokens)
        assertEquals(500L, record.latencyMs)
    }

    @Test
    fun `should handle multiple events correctly`() {
        // 1. Start
        val initialState = projector.initial()
        val startEvent = stored(
            payload = InferenceStartedEvent(
                requestId = InferenceRequestId("req1"),
                sessionId = SessionId("sess1"),
                stageId = StageId("stageA"),
                providerId = ProviderId("providerX"),
                promptArtifactId = TypeId("00".repeat(32)),
            ),
        )
        var currentState = projector.apply(initialState, startEvent)
        assertEquals(1, currentState.records.size)

        // 2. Complete
        val completeEvent = stored(
            payload = InferenceCompletedEvent(
                requestId = InferenceRequestId("req1"),
                sessionId = SessionId("sess1"),
                stageId = StageId("stageA"),
                providerId = ProviderId("providerX"),
                tokensUsed = TokenUsage(100, 100),
                latencyMs = 1000L,
                responseArtifactId = TypeId("00".repeat(32)),
            ),
        )
        currentState = projector.apply(currentState, completeEvent)
        assertEquals(1, currentState.records.size)
        assertEquals(InferenceStatus.COMPLETED, currentState.records.first().status)
        assertEquals(200, currentState.records.first().tokensUsed?.totalTokens)
    }
}
