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.events.events.ClarificationAnswer
import com.correx.core.events.events.ClarificationAnsweredEvent
import com.correx.core.events.events.ClarificationRequestedEvent
import com.correx.core.events.events.WorkflowCompletedEvent
import com.correx.core.events.execution.RetryPolicy
import com.correx.core.events.types.ArtifactId
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.InferenceRouter
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.cyclePolicyMissingValidator
import com.correx.testing.fixtures.context.ContextFixtures
import com.correx.testing.fixtures.inference.MockInferenceProvider
import com.correx.testing.fixtures.transitions.TransitionFixtures
import com.correx.testing.kernel.MockSessionEventReplayer
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.yield
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

/**
 * Integration test for the producer-exit clarification loop: a stage whose produced artifact
 * carries an unanswered `questions` array parks its run, surfaces a [ClarificationRequestedEvent],
 * and on operator answer re-runs the same stage with the answers in context — without consuming the
 * failure-retry budget (retry maxAttempts is 1 here). Bounded at three rounds.
 */
class ClarificationLoopTest {

    private val eventStore = InMemoryEventStore()
    private val sessionRepository = DefaultSessionRepository(MockSessionEventReplayer())
    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 repositories = OrchestratorRepositories(
        eventStore = eventStore,
        inferenceRepository = inferenceRepository,
        orchestrationRepository = orchestrationRepository,
        sessionRepository = sessionRepository,
        artifactRepository = LiveArtifactRepository(eventStore, DefaultArtifactReducer()),
        approvalRepository = approvalRepository,
    )

    private val analysisKind = ConfigArtifactKind(
        id = "analysis",
        schema = JsonSchema(type = "object", properties = emptyMap()),
        llmEmitted = true,
    )
    private val analysisId = ArtifactId("analysis")
    private val analystStage = StageId("analyst")

    /** Single stage that produces `analysis`, then transitions to the terminal `done`. */
    private fun graph() = WorkflowGraph(
        id = "clarify",
        stages = mapOf(
            analystStage to StageConfig(
                produces = listOf(TypedArtifactSlot(name = analysisId, kind = analysisKind)),
            ),
        ),
        transitions = setOf(
            TransitionEdge(
                id = TransitionId("analyst-to-done"),
                from = analystStage,
                to = StageId("done"),
                condition = { true },
            ),
        ),
        start = analystStage,
    )

    private val withQuestions =
        """{"summary":"build a UI","questions":[{"id":"stack","prompt":"Which stack?","options":["React","Vue"]}]}"""
    private val clean = """{"summary":"React chosen — no open questions"}"""

    /** Provider returns [responses] in order, clamping to the last once exhausted. */
    private fun orchestrator(responses: List<String>): DefaultSessionOrchestrator {
        var index = 0
        val router = object : InferenceRouter {
            override suspend fun route(
                stageId: StageId,
                requiredCapabilities: Set<ModelCapability>,
            ) = MockInferenceProvider(fixedResponse = responses[minOf(index++, responses.lastIndex)])
        }
        return DefaultSessionOrchestrator(
            repositories = repositories,
            engines = OrchestratorEngines(
                transitionResolver = TransitionFixtures.simpleResolver(),
                contextPackBuilder = ContextFixtures.simpleBuilder(),
                inferenceRouter = router,
                validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator())),
                approvalEngine = DefaultApprovalEngine(),
                riskAssessor = DefaultRiskAssessor(),
            ),
            retryCoordinator = DefaultRetryCoordinator(eventStore),
            artifactStore = NoopArtifactStore(),
            decisionJournalRepository = decisionJournalRepository,
        )
    }

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

    @Test
    fun `open questions park the stage, the answer re-runs it, then the workflow completes`():
        Unit = runBlocking {
        val sessionId = SessionId("clarify-happy")
        val orchestrator = orchestrator(listOf(withQuestions, clean))
        val runJob = launch { orchestrator.run(sessionId, graph(), config) }

        val request = withTimeout(5_000) {
            var r: ClarificationRequestedEvent? = null
            while (r == null) {
                r = eventStore.read(sessionId).firstNotNullOfOrNull { it.payload as? ClarificationRequestedEvent }
                if (r == null) yield()
            }
            r
        }
        assertEquals(analystStage, request.stageId)
        assertEquals("Which stack?", request.questions.single().prompt)

        orchestrator.submitClarification(
            sessionId, request.stageId, request.requestId,
            listOf(ClarificationAnswer(questionId = "stack", value = "React")),
        )
        runJob.join()

        val events = eventStore.read(sessionId)
        assertNotNull(
            events.firstNotNullOfOrNull { it.payload as? ClarificationAnsweredEvent },
            "Expected a ClarificationAnsweredEvent",
        )
        assertNotNull(
            events.find { it.payload is WorkflowCompletedEvent },
            "Workflow should complete after the clarification re-run produced a clean analysis",
        )
    }

    @Test
    fun `clarification rounds are capped so a stuck stage still proceeds`(): Unit = runBlocking {
        val sessionId = SessionId("clarify-cap")
        // Always returns questions: without a cap this would loop forever.
        val orchestrator = orchestrator(listOf(withQuestions))
        val runJob = launch { orchestrator.run(sessionId, graph(), config) }

        val handled = mutableSetOf<String>()
        withTimeout(10_000) {
            while (eventStore.read(sessionId).none { it.payload is WorkflowCompletedEvent }) {
                val pending = eventStore.read(sessionId)
                    .mapNotNull { it.payload as? ClarificationRequestedEvent }
                    .firstOrNull { it.requestId.value !in handled }
                if (pending != null) {
                    handled += pending.requestId.value
                    orchestrator.submitClarification(
                        sessionId, pending.stageId, pending.requestId,
                        listOf(ClarificationAnswer(questionId = "stack", value = "React")),
                    )
                }
                yield()
            }
        }
        runJob.join()

        val rounds = eventStore.read(sessionId).count { it.payload is ClarificationRequestedEvent }
        assertEquals(3, rounds, "Clarification must be capped at three rounds")
        assertNotNull(
            eventStore.read(sessionId).find { it.payload is WorkflowCompletedEvent },
            "Workflow should proceed once the clarification cap is reached",
        )
    }
}