package com.correx.testing.fixtures

import com.correx.core.context.model.ContextPack
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.EventId
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.GenerationConfig
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.TokenUsage
import com.correx.core.utils.TypeId
import com.correx.testing.fixtures.inference.MockInferenceProvider
import kotlinx.datetime.Clock
import java.util.*

@Suppress("LongParameterList")
object InferenceFixtures {

    fun fixedProvider(): InferenceProvider {
        return MockInferenceProvider(
            fixedResponse = "Fixed test response",
            artificialDelayMs = 0,
        )
    }

    fun fixedRouter(): InferenceRouter {
        return object : InferenceRouter {
            override suspend fun route(
                stageId: StageId,
                requiredCapabilities: Set<ModelCapability>,
            ): InferenceProvider {
                return fixedProvider()
            }
        }
    }

    fun failingRouter(): InferenceRouter {
        return object : InferenceRouter {
            override suspend fun route(
                stageId: StageId,
                requiredCapabilities: Set<ModelCapability>,
            ): InferenceProvider {
                return MockInferenceProvider(forcedFailure = "Simulated failure")
            }
        }
    }

    fun contextPack(): ContextPack {
        return ContextPack(
            id = ContextPackId("cp1"),
            sessionId = SessionId("s1"),
            stageId = StageId("stage1"),
            layers = emptyMap(),
            budgetUsed = 0,
            budgetLimit = 4096,
        )
    }

    fun generationConfig(): GenerationConfig {
        return GenerationConfig(
            temperature = 0.7,
            topP = 1.0,
            maxTokens = 2048,
        )
    }

    private const val PLACEHOLDER_HEX_BYTES = 32
    private val PLACEHOLDER_ARTIFACT_ID: ArtifactId = TypeId("00".repeat(PLACEHOLDER_HEX_BYTES))

    fun inferenceCompleted(
        requestId: InferenceRequestId,
        sessionId: SessionId,
        stageId: StageId,
        providerId: ProviderId,
        tokensUsed: TokenUsage,
        latencyMs: Long,
        responseArtifactId: ArtifactId = PLACEHOLDER_ARTIFACT_ID,
    ): InferenceCompletedEvent {
        return InferenceCompletedEvent(
            requestId = requestId,
            sessionId = sessionId,
            stageId = stageId,
            providerId = providerId,
            tokensUsed = tokensUsed,
            latencyMs = latencyMs,
            responseArtifactId = responseArtifactId,
        )
    }

    fun inferenceRequest(
        sessionId: SessionId,
        stageId: StageId,
        requestId: InferenceRequestId? = null,
    ): InferenceRequest {
        return InferenceRequest(
            requestId = requestId ?: InferenceRequestId(UUID.randomUUID().toString()),
            sessionId = sessionId,
            stageId = stageId,
            contextPack = contextPack(),
            generationConfig = generationConfig(),
        )
    }

    fun inferenceResponse(
        requestId: InferenceRequestId,
        text: String = "test response",
        tokensUsed: TokenUsage = TokenUsage(promptTokens = 100, completionTokens = 200),
        latencyMs: Long = 500,
    ): InferenceResponse {
        return InferenceResponse(
            requestId = requestId,
            text = text,
            finishReason = FinishReason.Stop,
            tokensUsed = tokensUsed,
            latencyMs = latencyMs,
        )
    }

    fun newInferenceCompletedEvent(
        requestId: InferenceRequestId,
        sessionId: SessionId,
        stageId: StageId,
        providerId: ProviderId,
        tokensUsed: TokenUsage,
        latencyMs: Long,
        responseArtifactId: ArtifactId = PLACEHOLDER_ARTIFACT_ID,
    ): NewEvent {
        return NewEvent(
            metadata = EventMetadata(
                eventId = EventId("inf-$requestId"),
                sessionId = sessionId,
                timestamp = Clock.System.now(),
                schemaVersion = 1,
                causationId = null,
                correlationId = null,
            ),
            payload = InferenceCompletedEvent(
                requestId = requestId,
                sessionId = sessionId,
                stageId = stageId,
                providerId = providerId,
                tokensUsed = tokensUsed,
                latencyMs = latencyMs,
                responseArtifactId = responseArtifactId,
            ),
        )
    }
}
