import com.correx.core.approvals.ApprovalOutcome
import com.correx.core.approvals.ApprovalProjector
import com.correx.core.approvals.ApprovalStatus
import com.correx.core.approvals.DefaultApprovalReducer
import com.correx.core.approvals.DefaultApprovalRepository
import com.correx.core.approvals.Tier
import com.correx.core.approvals.domain.ApprovalEngine
import com.correx.core.approvals.domain.DefaultApprovalEngine
import com.correx.core.approvals.model.ApprovalContext
import com.correx.core.approvals.model.ApprovalDecision
import com.correx.core.approvals.model.ApprovalGrant
import com.correx.core.approvals.model.ApprovalScopeIdentity
import com.correx.core.approvals.model.DomainApprovalRequest
import com.correx.core.artifacts.DefaultArtifactReducer
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.context.builder.DefaultContextPackBuilder
import com.correx.core.context.compression.CompressionStrategy
import com.correx.core.context.compression.ContextCompressor
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.ContextAssembledEvent
import com.correx.core.events.events.ContextTruncatedEvent
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.ToolInvocationRequestedEvent
import com.correx.core.events.events.WorkflowCompletedEvent
import com.correx.core.events.events.RetryAttemptedEvent
import com.correx.core.events.events.WorkflowFailedEvent
import com.correx.core.events.events.WorkflowStartedEvent
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.execution.RetryPolicy
import com.correx.core.events.types.ArtifactId
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.events.types.ToolInvocationId
import com.correx.core.events.types.TransitionId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.FinishReason
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRepository
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.DefaultInferenceRouter
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.InferenceState
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoEligibleProviderException
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.ProviderRegistry
import com.correx.core.inference.RoutingStrategy
import com.correx.core.inference.Token
import com.correx.core.inference.TokenUsage
import com.correx.core.inference.Tokenizer
import com.correx.core.inference.ToolCallFunction
import com.correx.core.inference.ToolCallRequest
import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer
import com.correx.core.journal.DecisionJournalProjector
import com.correx.core.journal.DefaultDecisionJournalReducer
import com.correx.core.journal.DefaultDecisionJournalRepository
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.ApprovalMode
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.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.transitions.evaluation.PromptResolver
import com.correx.core.transitions.resolution.DefaultTransitionResolver
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.utils.TypeId
import com.correx.core.validation.approval.ApprovalTrigger
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.context.ContextFixtures
import com.correx.testing.fixtures.cyclePolicyMissingValidator
import com.correx.testing.fixtures.inference.MockInferenceProvider
import com.correx.testing.fixtures.transitions.TransitionFixtures
import com.correx.testing.fixtures.transitions.TransitionFixtures.threeStageGraph
import com.correx.testing.kernel.MockSessionEventReplayer
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.yield
import kotlinx.datetime.Clock
import kotlinx.datetime.Instant
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
import java.util.concurrent.atomic.AtomicInteger

@Suppress("UnusedPrivateProperty")
class SessionOrchestratorIntegrationTest {
    private val eventStore = InMemoryEventStore()
    private val graph = TransitionFixtures.simpleGraph()
    private val transitionResolver = TransitionFixtures.simpleResolver()
    private val contextPackBuilder = ContextFixtures.simpleBuilder()
    private val inferenceRouter = InferenceFixtures.fixedRouter()
    private val validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator()))
    private val approvalEngine = DefaultApprovalEngine()
    private val retryCoordinator = DefaultRetryCoordinator(eventStore)
    private val sessionReplayer = MockSessionEventReplayer()
    private val sessionRepository = DefaultSessionRepository(sessionReplayer)
    val orchestrationReplayer = DefaultEventReplayer(
        eventStore,
        OrchestrationProjector(DefaultOrchestrationReducer()),
    )
    val orchestrationRepository = OrchestrationRepository(orchestrationReplayer)
    private val inferenceRepository = InferenceRepository(
        object : EventReplayer<InferenceState> {
            override fun rebuild(sessionId: SessionId) = InferenceState()
        },
    )
    private val riskAssessor = DefaultRiskAssessor()
    private val artifactStore = NoopArtifactStore()

    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 engines = OrchestratorEngines(
        transitionResolver = transitionResolver,
        contextPackBuilder = contextPackBuilder,
        inferenceRouter = inferenceRouter,
        validationPipeline = validationPipeline,
        approvalEngine = approvalEngine,
        riskAssessor = riskAssessor,
    )

    private val orchestrator = DefaultSessionOrchestrator(
        repositories = repositories,
        engines = engines,
        retryCoordinator = retryCoordinator,
        artifactStore = artifactStore,
        decisionJournalRepository = decisionJournalRepository,
    )

    @Test
    fun `successful workflow completes all stages`(): Unit = runBlocking {
        val sessionId = SessionId("s1")
        val config = OrchestrationConfig(
            retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0),
        )

        orchestrator.run(sessionId, graph, config)

        val events = eventStore.read(sessionId)

        val workflowStarted = events.find { it.payload is WorkflowStartedEvent }
        assertNotNull(workflowStarted)
        assertEquals(
            StageId("A"),
            (workflowStarted?.payload as? WorkflowStartedEvent)?.startStageId,
        )

        val workflowCompleted = events.find { it.payload is WorkflowCompletedEvent }
        assertNotNull(workflowCompleted)
        assertEquals(
            StageId("B"),
            (workflowCompleted?.payload as WorkflowCompletedEvent).terminalStageId,
        )
        assertEquals(1, (workflowCompleted.payload as WorkflowCompletedEvent).totalStages)
    }

    @Test
    fun `workflow fails with retry exhaustion`(): Unit = runBlocking {
        val sessionId = SessionId("s2")
        val config = OrchestrationConfig(
            retryPolicy = RetryPolicy(maxAttempts = 2, backoffMs = 0),
        )
        val graph = threeStageGraph()

        val failingOrchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines.copy(
                inferenceRouter = object : InferenceRouter {
                    override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>) =
                        MockInferenceProvider(forcedFailure = "Simulated failure")
                },
            ),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )
        failingOrchestrator.run(sessionId, graph, config)

        val events = eventStore.read(sessionId)
        val workflowFailed = events.find { it.payload is WorkflowFailedEvent }
        assertNotNull(workflowFailed)
        val failed = workflowFailed?.payload as WorkflowFailedEvent
        assertEquals("Simulated failure", failed.reason)
        assertTrue(failed.retryExhausted)
    }

    @Test
    fun `transient NoEligibleProviderException on a retry is tolerated, not a session-killing crash (#299)`(): Unit =
        runBlocking {
            val sessionId = SessionId("s2b")
            val config = OrchestrationConfig(
                retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0),
            )
            val graph = threeStageGraph()
            var routeCalls = 0

            val recoveringOrchestrator = DefaultSessionOrchestrator(
                repositories = repositories,
                engines = engines.copy(
                    inferenceRouter = object : InferenceRouter {
                        override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>): InferenceProvider {
                            routeCalls++
                            // First attempt hits the provider mid-crash — subsequent retries find it back.
                            if (routeCalls == 1) throw NoEligibleProviderException(stageId, requiredCapabilities)
                            return MockInferenceProvider()
                        }
                    },
                ),
                retryCoordinator = retryCoordinator,
                artifactStore = artifactStore,
                decisionJournalRepository = decisionJournalRepository,
            )
            recoveringOrchestrator.run(sessionId, graph, config)

            val events = eventStore.read(sessionId)
            // A crash would have skipped straight to an unhandled failure with no WorkflowCompletedEvent
            // and no orderly retry — assert the session survived and made it through via the retry path.
            assertTrue(routeCalls > 1) { "expected routing to be retried, only called $routeCalls time(s)" }
            assertNotNull(events.find { it.payload is WorkflowCompletedEvent })
            assertTrue(events.none { it.payload is WorkflowFailedEvent })
        }

    @Test
    fun `connection-level inference failure gates routing away from the dead provider immediately (#300)`(): Unit =
        runBlocking {
            val sessionId = SessionId("s2c")
            val config = OrchestrationConfig(
                retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0),
            )
            val graph = threeStageGraph()

            val dead = MockInferenceProvider(
                id = ProviderId("dead"),
                forcedFailure = "Failed to parse HTTP response: the server prematurely closed the connection",
            )
            val backup = MockInferenceProvider(id = ProviderId("backup"))
            // Router's own health probe (dead.healthCheck()) still reports Healthy — it hasn't
            // polled yet — so only the event-driven reportFailure() from the orchestrator's catch
            // block (not the periodic poll) can gate the retry away from `dead`.
            val registry = object : ProviderRegistry {
                override fun register(provider: com.correx.core.inference.InferenceProvider) = Unit
                override fun resolve(capability: ModelCapability) = listOf(dead, backup)
                override fun listAll() = listOf(dead, backup)
                override suspend fun healthCheckAll() = mapOf(dead.id to ProviderHealth.Healthy, backup.id to ProviderHealth.Healthy)
            }
            val router = DefaultInferenceRouter(
                registry = registry,
                strategy = RoutingStrategy { candidates, _ -> candidates.first() },
            )

            val recoveringOrchestrator = DefaultSessionOrchestrator(
                repositories = repositories,
                engines = engines.copy(inferenceRouter = router),
                retryCoordinator = retryCoordinator,
                artifactStore = artifactStore,
                decisionJournalRepository = decisionJournalRepository,
            )
            recoveringOrchestrator.run(sessionId, graph, config)

            val events = eventStore.read(sessionId)
            assertNotNull(events.find { it.payload is WorkflowCompletedEvent })
            assertEquals(1, dead.inferCallCount) { "dead provider should only be attempted once, then gated out" }
        }

    @Test
    fun `stage that matches no transition retries through the budget before failing`(): Unit = runBlocking {
        val sessionId = SessionId("s-no-transition-retry")
        // Stage A succeeds (no produces, no tools) but its only outgoing edge never matches, so the
        // resolver returns Stay. That must consume the retry budget, not fail the run on first sight.
        val graph = WorkflowGraph(
            id = "no-transition-retry-test",
            stages = mapOf(StageId("A") to StageConfig()),
            transitions = setOf(
                TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { false }),
            ),
            start = StageId("A"),
        )
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 2, backoffMs = 0))
        // The default fixture evaluator returns true unconditionally; use one that honours the
        // condition so the { false } edge actually produces a Stay.
        val orchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines.copy(
                transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) },
            ),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        orchestrator.run(sessionId, graph, config)

        val events = eventStore.read(sessionId)
        val retries = events.count { it.payload is RetryAttemptedEvent }
        assertEquals(2, retries, "expected the no-transition outcome to exhaust the retry budget")
        val failed = events.find { it.payload is WorkflowFailedEvent }?.payload as? WorkflowFailedEvent
        assertNotNull(failed)
        assertTrue(failed!!.retryExhausted, "expected retryExhausted=true once the budget is spent")
        assertTrue(
            failed.reason.contains("no transition condition matched"),
            "unexpected failure reason: ${failed.reason}",
        )
    }

    @Test
    fun `workflow pauses for approval when required`(): Unit = runBlocking {
        val sessionId = SessionId("s3")
        val config = OrchestrationConfig(
            retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0),
        )
        val graph = threeStageGraph()
        val approvingPipeline = ValidationPipeline(
            validators = listOf(cyclePolicyMissingValidator()),
            approvalTrigger = ApprovalTrigger(),
        )
        val orchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines.copy(validationPipeline = approvingPipeline),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        val runJob = launch { orchestrator.run(sessionId, graph, config) }

        withTimeout(5_000) {
            while (eventStore.read(sessionId).none { it.payload is OrchestrationPausedEvent }) {
                yield()
            }
        }

        val events = eventStore.read(sessionId)
        val orPause = events.find { it.payload is OrchestrationPausedEvent }
        assertNotNull(orPause)
        val paused = orPause?.let { it.payload as? OrchestrationPausedEvent }
        assertEquals("APPROVAL_PENDING", paused?.reason)

        runJob.cancel()
        runJob.join()
    }

    @Test
    fun `workflow resumes after approval`(): Unit = runBlocking {
        val sessionId = SessionId("s4")
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0))
        // Single-stage graph: enterStage(A) triggers exactly one approval; after resume,
        // the always-true transition moves to the terminal "done" and the workflow completes.
        val graph = WorkflowGraph(
            id = "workflow-test",
            stages = mapOf(StageId("A") to StageConfig()),
            transitions = setOf(
                TransitionEdge(
                    id = com.correx.core.events.types.TransitionId("t1"),
                    from = StageId("A"),
                    to = StageId("done"),
                    condition = { true },
                ),
            ),
            start = StageId("A"),
        )
        val approvingPipeline = ValidationPipeline(
            validators = listOf(cyclePolicyMissingValidator()),
            approvalTrigger = ApprovalTrigger(),
        )
        val approvalOrchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines.copy(validationPipeline = approvingPipeline),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        val runJob = launch { approvalOrchestrator.run(sessionId, graph, config) }

        val requestId = withTimeout(5_000) {
            var id: com.correx.core.events.types.ApprovalRequestId? = null
            while (id == null) {
                id = eventStore.read(sessionId)
                    .firstNotNullOfOrNull { it.payload as? ApprovalRequestedEvent }
                    ?.requestId
                if (id == null) yield()
            }
            id
        }

        val identity = ApprovalScopeIdentity(sessionId = sessionId, stageId = null, projectId = null)
        val context = ApprovalContext(identity = identity, mode = ApprovalMode.PROMPT)
        val decision = ApprovalDecision(
            id = null,
            requestId = requestId,
            outcome = ApprovalOutcome.APPROVED,
            state = ApprovalStatus.COMPLETED,
            tier = Tier.T2,
            contextSnapshot = context,
            resolutionTimestamp = Clock.System.now(),
            reason = null,
        )
        approvalOrchestrator.submitApprovalDecision(requestId, decision)

        runJob.join()

        val events = eventStore.read(sessionId)
        assertEquals(1, events.count { it.payload is OrchestrationResumedEvent })
        assertNotNull(events.find { it.payload is WorkflowCompletedEvent })
    }

    @Test
    fun `workflow with empty graph fails immediately`(): Unit = runBlocking {
        val emptyGraph = WorkflowGraph(
            id = "workflow-test",
            stages = mapOf(StageId("start") to StageConfig()),
            transitions = emptySet(),
            start = StageId("start"),
        )
        val config = OrchestrationConfig(
            retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0),
        )

        val orchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines,
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        val sessionId = SessionId("s6")
        orchestrator.run(sessionId, emptyGraph, config)

        val events = eventStore.read(sessionId)
        val workflowFailed = events.find { it.payload is WorkflowFailedEvent }
        assertNotNull(workflowFailed)
        val failed = workflowFailed?.payload as WorkflowFailedEvent
        assertNotNull(failed.reason)
    }

    @Test
    fun `context budget overflow on the workflow path emits ContextTruncatedEvent`(): Unit = runBlocking {
        // A compressor that drops everything forces the builder to report dropped entries; the
        // orchestrator must surface that as a ContextTruncatedEvent (mirroring the router path).
        val droppingBuilder = DefaultContextPackBuilder(
            object : ContextCompressor {
                override fun compress(
                    entries: List<ContextEntry>,
                    budget: TokenBudget,
                    strategy: CompressionStrategy,
                ): List<ContextEntry> = emptyList()
            },
        )
        val truncGraph = WorkflowGraph(
            id = "trunc-test",
            stages = mapOf(
                StageId("A") to StageConfig(metadata = mapOf("systemPrompt" to "you are a helpful assistant")),
            ),
            transitions = setOf(
                TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { true }),
            ),
            start = StageId("A"),
        )
        val truncOrchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            // PromptResolver { it } makes the systemPrompt metadata resolve to non-blank text,
            // so there is a compressible entry for the dropping compressor to discard.
            engines = engines.copy(contextPackBuilder = droppingBuilder, promptResolver = PromptResolver { it }),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )
        val sessionId = SessionId("s-trunc")
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0))

        // Seed a recorded repo map so the stage context carries an OPTIONAL entry the dropping
        // compressor can discard: required entries (like the system prompt) are pinned by the
        // bucket split and never reach the compressor.
        eventStore.append(
            com.correx.core.events.events.NewEvent(
                com.correx.core.events.events.EventMetadata(
                    com.correx.core.events.types.EventId("rm-trunc"),
                    sessionId,
                    kotlinx.datetime.Instant.parse("2026-06-04T00:00:00Z"),
                    1,
                    null,
                    null,
                ),
                com.correx.core.events.events.RepoMapComputedEvent(
                    sessionId,
                    "/repo",
                    listOf(com.correx.core.events.events.RepoMapEntry("src/Main.kt", 0.9)),
                    kotlinx.datetime.Instant.parse("2026-06-04T00:00:00Z"),
                ),
            ),
        )

        truncOrchestrator.run(sessionId, truncGraph, config)

        val truncated = eventStore.read(sessionId).mapNotNull { it.payload as? ContextTruncatedEvent }
        assertTrue(truncated.isNotEmpty(), "expected a ContextTruncatedEvent when the budget drops entries")
        assertEquals("A", truncated.first().turnId)
        assertTrue(truncated.first().entriesDropped >= 1)
    }

    @Test
    fun `stage context build emits ContextAssembledEvent with a manifest (#307)`(): Unit = runBlocking {
        val sessionId = SessionId("s-manifest")
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0))

        orchestrator.run(sessionId, graph, config)

        val assembled = eventStore.read(sessionId).mapNotNull { it.payload as? ContextAssembledEvent }
        assertTrue(assembled.isNotEmpty(), "expected a ContextAssembledEvent per stage context build")
        val first = assembled.first()
        assertEquals(StageId("A"), first.stageId)
        assertTrue(first.entries.isNotEmpty(), "manifest should list the entries injected into the stage")
        // Manifest-only: identifying fields present, never the content (content lives in CAS/replay).
        assertTrue(first.entries.all { it.sourceType.isNotBlank() && it.layer.isNotBlank() && it.role.isNotBlank() })
    }

    @Test
    fun `artifactStore put is called for prompt and response and ids appear on inference events`(): Unit = runBlocking {
        val recordingStore = RecordingArtifactStore()
        val sessionId = SessionId("s-artifact")
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0))
        val recordingOrchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines,
            retryCoordinator = retryCoordinator,
            artifactStore = recordingStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        recordingOrchestrator.run(sessionId, graph, config)

        assertTrue(recordingStore.puts.size >= 2, "Expected at least 2 puts, got ${recordingStore.puts.size}")

        val events = eventStore.read(sessionId)
        val started = events.firstNotNullOf { it.payload as? InferenceStartedEvent }
        val completed = events.firstNotNullOf { it.payload as? InferenceCompletedEvent }

        assertEquals(recordingStore.idForIndex(0), started.promptArtifactId)
        assertEquals(recordingStore.idForIndex(1), completed.responseArtifactId)
    }

    @Test
    fun `run throws IllegalArgumentException when transition targets undeclared stage`(): Unit = runBlocking {
        val a = StageId("A")
        val b = StageId("B")
        val ghost = StageId("ghost")
        val z = StageId("Z")
        // ghost is missing from stages map but has an outgoing transition so it is non-terminal
        val brokenGraph = WorkflowGraph(
            id = "workflow-test",
            stages = mapOf(a to StageConfig(), b to StageConfig()),
            transitions = setOf(
                TransitionEdge(
                    id = com.correx.core.events.types.TransitionId("t1"), from = a, to = b,
                    condition = { true },
                ),
                TransitionEdge(
                    id = com.correx.core.events.types.TransitionId("t2"), from = b, to = ghost,
                    condition = { true },
                ),
                TransitionEdge(
                    id = com.correx.core.events.types.TransitionId("t3"), from = ghost, to = z,
                    condition = { true },
                ),
            ),
            start = a,
        )
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0))
        val sessionId = SessionId("s7")

        val ex = org.junit.jupiter.api.Assertions.assertThrows(IllegalArgumentException::class.java) {
            runBlocking { orchestrator.run(sessionId, brokenGraph, config) }
        }
        assertTrue(ex.message?.contains("ghost") == true, "Expected message to mention 'ghost', got: ${ex.message}")
    }

    @Test
    fun `WorkflowStartedEvent carries retryPolicy from config`(): Unit = runBlocking {
        val sessionId = SessionId("s8")
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 5, backoffMs = 0))

        orchestrator.run(sessionId, graph, config)

        val events = eventStore.read(sessionId)
        val started = events.first { it.payload is WorkflowStartedEvent }.payload as WorkflowStartedEvent
        assertEquals(5, started.retryPolicy?.maxAttempts)
    }

    @Test
    fun `OrchestrationState retryPolicy reflects config after workflow starts`(): Unit = runBlocking {
        val sessionId = SessionId("s9")
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 7, backoffMs = 0))

        orchestrator.run(sessionId, graph, config)

        val state = orchestrationRepository.getState(sessionId)
        assertEquals(7, state.retryPolicy?.maxAttempts)
    }

    @Test
    fun `stage with allowedTools but no tool invocation fails liveness check`(): Unit = runBlocking {
        val sessionId = SessionId("s-liveness-neg")
        // Stage offers a tool but the stubbed LLM immediately returns Stop (stage_complete
        // shortcut) without ever calling it. The new gate must reject this vacuous pass.
        val livenessGraph = WorkflowGraph(
            id = "liveness-test",
            stages = mapOf(
                StageId("A") to StageConfig(allowedTools = setOf("my_tool")),
            ),
            transitions = setOf(
                TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { true }),
            ),
            start = StageId("A"),
        )
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 2, backoffMs = 0))
        // Default MockInferenceProvider returns FinishReason.Stop with no tool calls.
        val livenessOrchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines,
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        livenessOrchestrator.run(sessionId, livenessGraph, config)

        val events = eventStore.read(sessionId)
        val failed = events.find { it.payload is WorkflowFailedEvent }
        assertNotNull(failed, "expected WorkflowFailedEvent when tool was offered but never called")
        val payload = failed?.payload as WorkflowFailedEvent
        assertTrue(payload.retryExhausted, "expected retryExhausted=true")
        assertTrue(
            payload.reason.contains("declared no artifacts and ran no tools"),
            "unexpected failure reason: ${payload.reason}",
        )
    }

    @Test
    fun `stage with allowedTools passes liveness check after tool is invoked`(): Unit = runBlocking {
        val sessionId = SessionId("s-liveness-pos")
        val livenessGraph = WorkflowGraph(
            id = "liveness-pos-test",
            stages = mapOf(
                StageId("A") to StageConfig(allowedTools = setOf("my_tool")),
            ),
            transitions = setOf(
                TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { true }),
            ),
            start = StageId("A"),
        )
        val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0))

        // Round 0: return a tool call for "my_tool"; round 1: call stage_complete to finish.
        val callCount = AtomicInteger(0)
        val sequencedProvider = object : InferenceProvider {
            override val id: ProviderId = ProviderId("seq-mock")
            override val name: String = "Sequenced Mock"
            override val tokenizer: Tokenizer = object : Tokenizer {
                override suspend fun tokenize(text: String): List<Token> = emptyList()
                override suspend fun countTokens(text: String) = text.length / 4
            }
            override suspend fun infer(request: InferenceRequest): InferenceResponse {
                val round = callCount.getAndIncrement()
                return if (round == 0) {
                    InferenceResponse(
                        requestId = request.requestId,
                        text = "",
                        finishReason = FinishReason.ToolCall,
                        tokensUsed = TokenUsage(promptTokens = 10, completionTokens = 5),
                        latencyMs = 0,
                        toolCalls = listOf(
                            ToolCallRequest(
                                id = "call-1",
                                function = ToolCallFunction(name = "my_tool", arguments = "{}"),
                            ),
                        ),
                    )
                } else {
                    InferenceResponse(
                        requestId = request.requestId,
                        text = "",
                        finishReason = FinishReason.ToolCall,
                        tokensUsed = TokenUsage(promptTokens = 10, completionTokens = 5),
                        latencyMs = 0,
                        toolCalls = listOf(
                            ToolCallRequest(
                                id = "call-2",
                                function = ToolCallFunction(name = "stage_complete", arguments = "{}"),
                            ),
                        ),
                    )
                }
            }
            override suspend fun healthCheck() = ProviderHealth.Healthy
            override fun capabilities() = setOf(CapabilityScore(ModelCapability.General, 1.0))
        }

        // Auto-approving engine so T2 tool calls are not blocked in tests.
        val autoApproveEngine = object : ApprovalEngine {
            override fun evaluate(
                request: DomainApprovalRequest,
                context: ApprovalContext,
                grants: List<ApprovalGrant>,
                now: Instant,
            ) = ApprovalDecision(
                id = null,
                requestId = request.id,
                outcome = ApprovalOutcome.AUTO_APPROVED,
                state = ApprovalStatus.COMPLETED,
                tier = request.tier,
                contextSnapshot = context,
                resolutionTimestamp = now,
                reason = "test-auto-approve",
            )
        }

        // Tool executor that succeeds for any tool call.
        val succeedingExecutor = object : ToolExecutor {
            override suspend fun execute(request: ToolRequest) = ToolResult.Success(
                invocationId = request.invocationId,
                output = "ok",
            )
        }

        val livenessOrchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines.copy(
                inferenceRouter = object : InferenceRouter {
                    override suspend fun route(
                        stageId: StageId,
                        requiredCapabilities: Set<ModelCapability>,
                    ): InferenceProvider = sequencedProvider
                },
                approvalEngine = autoApproveEngine,
                toolExecutor = succeedingExecutor,
            ),
            retryCoordinator = retryCoordinator,
            artifactStore = artifactStore,
            decisionJournalRepository = decisionJournalRepository,
        )

        livenessOrchestrator.run(sessionId, livenessGraph, config)

        val events = eventStore.read(sessionId)
        assertNotNull(
            events.find { it.payload is WorkflowCompletedEvent },
            "expected WorkflowCompletedEvent when tool was invoked before stage_complete",
        )
        assertTrue(
            events.any { (it.payload as? ToolInvocationRequestedEvent)?.stageId == StageId("A") },
            "expected ToolInvocationRequestedEvent for stage A",
        )
    }

}

private class RecordingArtifactStore : ArtifactStore {
    val puts: MutableList<ByteArray> = mutableListOf()
    private val byId: MutableMap<ArtifactId, ByteArray> = mutableMapOf()

    fun idForIndex(i: Int): ArtifactId = TypeId("0".repeat(63) + i.toString(16))

    override suspend fun put(bytes: ByteArray): ArtifactId {
        val id = idForIndex(puts.size)
        puts.add(bytes)
        byId[id] = bytes
        return id
    }

    override suspend fun get(id: ArtifactId): ByteArray? = byId[id]

    override suspend fun flushBefore(commit: suspend () -> Unit) {
        commit()
    }
}
