import com.correx.core.approvals.ApprovalProjector
import com.correx.core.approvals.DefaultApprovalReducer
import com.correx.core.approvals.DefaultApprovalRepository
import com.correx.core.approvals.domain.DefaultApprovalEngine
import com.correx.core.artifacts.DefaultArtifactReducer
import com.correx.core.artifacts.kind.ConfigArtifactKind
import com.correx.core.artifacts.kind.JsonSchema
import com.correx.core.artifacts.kind.TypedArtifactSlot
import com.correx.core.context.builder.ContextPackBuilder
import com.correx.core.context.model.CompressionMetadata
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.execution.RetryPolicy
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.TransitionId
import com.correx.core.inference.InferenceRepository
import com.correx.core.inference.InferenceState
import com.correx.core.inference.ModelCapability
import com.correx.core.journal.DecisionJournalProjector
import com.correx.core.journal.DefaultDecisionJournalReducer
import com.correx.core.journal.DefaultDecisionJournalRepository
import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer
import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator
import com.correx.core.kernel.orchestration.OrchestrationConfig
import com.correx.core.kernel.orchestration.OrchestrationProjector
import com.correx.core.kernel.orchestration.OrchestrationRepository
import com.correx.core.kernel.orchestration.OrchestratorEngines
import com.correx.core.kernel.orchestration.OrchestratorRepositories
import com.correx.core.kernel.retry.DefaultRetryCoordinator
import com.correx.core.risk.DefaultRiskAssessor
import com.correx.core.sessions.DefaultSessionRepository
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.sessions.projections.replay.EventReplayer
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.TransitionEdge
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.core.validation.pipeline.ValidationPipeline
import com.correx.infrastructure.persistence.InMemoryEventStore
import com.correx.infrastructure.persistence.artifact.LiveArtifactRepository
import com.correx.testing.contracts.fixtures.artifactstore.NoopArtifactStore
import com.correx.testing.fixtures.InferenceFixtures
import com.correx.testing.fixtures.WorkflowFixtures
import com.correx.testing.fixtures.cyclePolicyMissingValidator
import com.correx.testing.fixtures.inference.MockInferenceProvider
import com.correx.testing.fixtures.transitions.TransitionFixtures
import com.correx.testing.kernel.MockSessionEventReplayer
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ArtifactCreatedEvent
import com.correx.core.events.events.ArtifactValidatedEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.ToolInvocationRequestedEvent
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.EventId
import com.correx.core.events.types.ToolInvocationId
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

/**
 * Integration tests verifying that a stage's declared [StageConfig.needs] artifact ids are
 * injected into its context pack as "neededArtifact" entries, sourced from artifactContentCache.
 */
class NeedsArtifactInjectionTest {

    /**
     * A capturing [ContextPackBuilder] that records every entry list it sees so tests
     * can assert on what was assembled for each stage.
     */
    private val capturedEntries = mutableListOf<ContextEntry>()

    private val capturingBuilder = object : ContextPackBuilder {
        override suspend fun build(
            id: ContextPackId,
            sessionId: SessionId,
            stageId: StageId,
            entries: List<ContextEntry>,
            budget: TokenBudget,
        ): ContextPack {
            capturedEntries.addAll(entries)
            return ContextPack(
                id = id,
                sessionId = sessionId,
                stageId = stageId,
                layers = emptyMap(),
                budgetUsed = 0,
                budgetLimit = budget.limit,
            )
        }
    }

    private val eventStore = InMemoryEventStore()
    private val sessionReplayer = MockSessionEventReplayer()
    private val sessionRepository = DefaultSessionRepository(sessionReplayer)
    private val orchestrationRepository = OrchestrationRepository(
        DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())),
    )
    private val inferenceRepository = InferenceRepository(
        object : EventReplayer<InferenceState> {
            override fun rebuild(sessionId: SessionId) = InferenceState()
        },
    )
    private val approvalRepository = DefaultApprovalRepository(
        DefaultEventReplayer(eventStore, ApprovalProjector(DefaultApprovalReducer())),
    )
    private val decisionJournalRepository = DefaultDecisionJournalRepository(
        DefaultEventReplayer(eventStore, DecisionJournalProjector(DefaultDecisionJournalReducer())),
    )
    private val artifactStore = NoopArtifactStore()
    private val retryCoordinator = DefaultRetryCoordinator(eventStore)

    private val repositories = OrchestratorRepositories(
        eventStore = eventStore,
        inferenceRepository = inferenceRepository,
        orchestrationRepository = orchestrationRepository,
        sessionRepository = sessionRepository,
        artifactRepository = LiveArtifactRepository(eventStore, DefaultArtifactReducer()),
        approvalRepository = approvalRepository,
    )

    private fun buildOrchestrator(inferenceRouter: com.correx.core.inference.InferenceRouter) =
        DefaultSessionOrchestrator(
            repositories = repositories,
            engines = OrchestratorEngines(
                transitionResolver = TransitionFixtures.simpleResolver(),
                contextPackBuilder = capturingBuilder,
                inferenceRouter = inferenceRouter,
                validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator())),
                approvalEngine = DefaultApprovalEngine(),
                riskAssessor = DefaultRiskAssessor(),
            ),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

    private val defaultConfig = OrchestrationConfig(
        retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0),
    )

    /**
     * Stage A produces an llm-emitted artifact ("analysis") so the orchestrator stores the
     * inference response text in artifactContentCache. Stage B declares needs=[analysis] and
     * should receive a "neededArtifact" context entry with that content.
     */
    @Test
    fun `stage with needs receives artifact content in context entries`(): Unit = runBlocking {
        val sessionId = SessionId("needs-test-1")
        val stageA = StageId("A")
        val stageB = StageId("B")
        val artifactId = ArtifactId("analysis")
        val artifactJson = """{"summary":"important findings"}"""

        val analysisKind = ConfigArtifactKind(
            id = "analysis",
            schema = JsonSchema(type = "object", properties = emptyMap()),
            llmEmitted = true,
        )

        val stageC = StageId("C")
        val graph = WorkflowGraph(
            id = "needs-injection-test",
            stages = mapOf(
                stageA to StageConfig(
                    produces = listOf(TypedArtifactSlot(name = artifactId, kind = analysisKind)),
                ),
                stageB to StageConfig(
                    needs = setOf(artifactId),
                ),
                stageC to StageConfig(),
            ),
            transitions = setOf(
                TransitionEdge(id = TransitionId("t1"), from = stageA, to = stageB, condition = { true }),
                TransitionEdge(id = TransitionId("t2"), from = stageB, to = stageC, condition = { true }),
            ),
            start = stageA,
        )

        // Stage A returns the artifact JSON so it gets cached; stages B and C return generic text.
        val router = object : com.correx.core.inference.InferenceRouter {
            private var callCount = 0
            override suspend fun route(
                stageId: StageId,
                requiredCapabilities: Set<ModelCapability>,
            ) = MockInferenceProvider(
                fixedResponse = if (callCount++ == 0) artifactJson else "stage response",
            )
        }

        buildOrchestrator(router).run(sessionId, graph, defaultConfig)

        val needsEntries = capturedEntries.filter { it.sourceType == "neededArtifact" }
        assertTrue(needsEntries.isNotEmpty(), "Expected at least one neededArtifact entry")
        val analysisEntry = needsEntries.find { it.sourceId == "analysis" }
        assertFalse(analysisEntry == null, "Expected a neededArtifact entry with sourceId=analysis")
        assertTrue(
            analysisEntry!!.content.contains(artifactJson),
            "Entry content should include artifact JSON, got: ${analysisEntry.content}",
        )
    }

    /**
     * Slice 6: a file written by an EARLIER stage this session must surface to a later stage as a
     * deterministic retrieval hit (a "relevantFiles" entry naming the path), even with no embedder —
     * the session-start repo map can't contain it and semantic retrieval never indexed it.
     */
    @Test
    fun `file written by an earlier stage surfaces to a later stage as a relevant file`(): Unit = runBlocking {
        val sessionId = SessionId("write-overlay-1")
        val stageA = StageId("A")
        val stageB = StageId("B")
        val inv = ToolInvocationId("inv-A-1")
        val siblingPath = "src/Sibling.kt"

        // Seed stage A's write into the log before the run; sessionWrittenHits projects it for stage B.
        listOf(
            ToolInvocationRequestedEvent(
                invocationId = inv, sessionId = sessionId, stageId = stageA,
                toolName = "file_write", tier = Tier.T2,
                request = ToolRequest(
                    invocationId = inv, sessionId = sessionId, stageId = stageA,
                    toolName = "file_write", parameters = mapOf("path" to siblingPath),
                ),
            ),
            FileWrittenEvent(
                invocationId = inv, sessionId = sessionId, path = siblingPath,
                postImageHash = "deadbeef", preExisted = false, timestampMs = 1,
            ),
        ).forEach { payload ->
            eventStore.append(
                NewEvent(
                    metadata = EventMetadata(
                        eventId = EventId(java.util.UUID.randomUUID().toString()),
                        sessionId = sessionId, timestamp = Clock.System.now(),
                        schemaVersion = 1, causationId = null, correlationId = null,
                    ),
                    payload = payload,
                ),
            )
        }

        // stageC makes stageB non-terminal so stageB actually gets a context pack built (terminal
        // stages are entered as end-markers without assembling context).
        val stageC = StageId("C")
        val graph = WorkflowGraph(
            id = "write-overlay-graph",
            stages = mapOf(stageA to StageConfig(), stageB to StageConfig(), stageC to StageConfig()),
            transitions = setOf(
                TransitionEdge(id = TransitionId("t1"), from = stageA, to = stageB, condition = { true }),
                TransitionEdge(id = TransitionId("t2"), from = stageB, to = stageC, condition = { true }),
            ),
            start = stageA,
        )
        buildOrchestrator(InferenceFixtures.fixedRouter()).run(sessionId, graph, defaultConfig)

        val relevant = capturedEntries.filter { it.sourceType == "relevantFiles" }
        assertTrue(
            relevant.any { it.content.contains(siblingPath) },
            "Expected a relevantFiles entry naming $siblingPath, got: ${relevant.map { it.content }}",
        )
    }

    @Test
    fun `stage with empty needs adds no neededArtifact entries`(): Unit = runBlocking {
        val sessionId = SessionId("needs-test-2")
        val graph = WorkflowFixtures.simpleGraph()

        buildOrchestrator(InferenceFixtures.fixedRouter()).run(sessionId, graph, defaultConfig)

        val needsEntries = capturedEntries.filter { it.sourceType == "neededArtifact" }
        assertTrue(needsEntries.isEmpty(), "No neededArtifact entries expected for stages with empty needs")
    }

    @Test
    fun `missing cache content for declared need is skipped without crash`(): Unit = runBlocking {
        val sessionId = SessionId("needs-test-3")
        val stageA = StageId("A")
        val stageB = StageId("B")
        val missingArtifactId = ArtifactId("missing_artifact")

        val graph = WorkflowGraph(
            id = "missing-needs-test",
            stages = mapOf(
                stageA to StageConfig(),
                stageB to StageConfig(needs = setOf(missingArtifactId)),
            ),
            transitions = setOf(
                TransitionEdge(id = TransitionId("t1"), from = stageA, to = stageB, condition = { true }),
            ),
            start = stageA,
        )

        // No llm-emitted artifact in stageA → cache never populated for missingArtifactId
        buildOrchestrator(InferenceFixtures.fixedRouter()).run(sessionId, graph, defaultConfig)

        val needsEntries = capturedEntries.filter { it.sourceType == "neededArtifact" }
        assertTrue(needsEntries.isEmpty(), "Missing cache content should produce no neededArtifact entries")
    }

    /**
     * Regression test for the phantom-artifact bug: when an llm-emitted slot's inference
     * response is blank, emitLlmArtifacts must NOT emit ArtifactCreatedEvent or
     * ArtifactValidatedEvent. Only a slot with non-blank cached content should produce events.
     */
    @Test
    fun `blank inference output does not emit artifact events for llm-emitted slot`(): Unit = runBlocking {
        val sessionId = SessionId("phantom-artifact-test-1")
        val stageA = StageId("A")
        val stageB = StageId("B")
        val blankArtifactId = ArtifactId("blank_output")

        val blankKind = ConfigArtifactKind(
            id = "blank_output",
            schema = JsonSchema(type = "object", properties = emptyMap()),
            llmEmitted = true,
        )

        val graph = WorkflowGraph(
            id = "phantom-artifact-graph",
            stages = mapOf(
                stageA to StageConfig(
                    produces = listOf(TypedArtifactSlot(name = blankArtifactId, kind = blankKind)),
                ),
                stageB to StageConfig(),
            ),
            transitions = setOf(
                TransitionEdge(id = TransitionId("t1"), from = stageA, to = stageB, condition = { true }),
            ),
            start = stageA,
        )

        // Stage A returns blank/whitespace — content cache will not be populated
        val router = object : com.correx.core.inference.InferenceRouter {
            override suspend fun route(
                stageId: StageId,
                requiredCapabilities: Set<ModelCapability>,
            ) = MockInferenceProvider(fixedResponse = "   ")
        }

        buildOrchestrator(router).run(sessionId, graph, defaultConfig)

        val events = eventStore.read(sessionId)
        val createdForBlank = events.filter { e ->
            (e.payload as? ArtifactCreatedEvent)?.artifactId == blankArtifactId
        }
        val validatedForBlank = events.filter { e ->
            (e.payload as? ArtifactValidatedEvent)?.artifactId == blankArtifactId
        }

        assertTrue(
            createdForBlank.isEmpty(),
            "ArtifactCreatedEvent must not be emitted when inference output is blank",
        )
        assertTrue(
            validatedForBlank.isEmpty(),
            "ArtifactValidatedEvent must not be emitted when inference output is blank",
        )
    }
}
