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.context.builder.ContextPackBuilder
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.events.EventMetadata
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.PreemptRedirectBlockedEvent
import com.correx.core.events.events.PreemptRedirectEvent
import com.correx.core.events.events.TransitionExecutedEvent
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.EventId
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.execution.WorkflowResult
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.inference.MockInferenceProvider
import com.correx.testing.fixtures.transitions.TransitionFixtures
import com.correx.testing.kernel.MockSessionEventReplayer
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import java.util.UUID

/**
 * Freestyle graph re-routing — the deterministic override half. A recorded, unconsumed
 * [PreemptRedirectEvent] overrides the resolver's edge (jump), unless the target is unknown or its
 * `needs` aren't satisfied (block). Exercises the override directly by emitting the confirmed
 * redirect event (the front-half's LLM-proposal + approval-confirm is out of scope here).
 */
class GraphRerouteTest {

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

    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 router = object : InferenceRouter {
        override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>) =
            MockInferenceProvider(fixedResponse = "done")
    }

    private fun orchestrator() = DefaultSessionOrchestrator(
        repositories = repositories,
        engines = OrchestratorEngines(
            transitionResolver = TransitionFixtures.simpleResolver(),
            contextPackBuilder = passthroughBuilder,
            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))

    // A → B → C → D, D terminal. All edges always-true, so the unredirected path is A→B→C→D.
    private fun linearGraph(cNeeds: Set<ArtifactId> = emptySet()): WorkflowGraph {
        val a = StageId("A"); val b = StageId("B"); val c = StageId("C"); val d = StageId("D")
        return WorkflowGraph(
            id = "reroute-test",
            stages = mapOf(
                a to StageConfig(), b to StageConfig(),
                c to StageConfig(needs = cNeeds), d to StageConfig(),
            ),
            transitions = setOf(
                TransitionEdge(id = TransitionId("ab"), from = a, to = b, condition = { true }),
                TransitionEdge(id = TransitionId("bc"), from = b, to = c, condition = { true }),
                TransitionEdge(id = TransitionId("cd"), from = c, to = d, condition = { true }),
            ),
            start = a,
        )
    }

    private suspend fun append(sessionId: SessionId, payload: EventPayload) {
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = payload,
            ),
        )
    }

    private fun transitions(sessionId: SessionId) =
        eventStore.read(sessionId).mapNotNull { it.payload as? TransitionExecutedEvent }

    @Test
    fun `confirmed redirect overrides the resolver edge and is consumed once`(): Unit = runBlocking {
        val sessionId = SessionId("reroute-jump")
        append(
            sessionId,
            PreemptRedirectEvent(
                redirectId = "r1",
                sessionId = sessionId,
                fromStageId = StageId("A"),
                toStageId = StageId("C"),
                timestampMs = 1L,
            ),
        )

        val result = orchestrator().run(sessionId, linearGraph(), config)

        val ts = transitions(sessionId)
        // A jumped straight to C — B was skipped entirely.
        assertTrue(ts.none { it.to == StageId("B") }, "B should be skipped, got: ${ts.map { it.to.value }}")
        val redirectHop = ts.single { it.redirectId == "r1" }
        assertEquals(StageId("C"), redirectHop.to)
        assertEquals(StageId("A"), redirectHop.from)
        // Consumed exactly once: C advances normally to the terminal D and the run completes,
        // rather than re-applying the redirect to C forever.
        assertTrue(result is WorkflowResult.Completed, "expected completion, got $result")
        assertEquals(StageId("D"), (result as WorkflowResult.Completed).terminalStageId)
    }

    @Test
    fun `redirect to a target with unsatisfied needs is blocked, normal edge taken`(): Unit = runBlocking {
        val sessionId = SessionId("reroute-block-needs")
        append(
            sessionId,
            PreemptRedirectEvent(
                redirectId = "r2",
                sessionId = sessionId,
                fromStageId = StageId("A"),
                toStageId = StageId("C"),
                timestampMs = 1L,
            ),
        )

        // C needs an artifact that is never produced/cached.
        orchestrator().run(sessionId, linearGraph(cNeeds = setOf(ArtifactId("design"))), config)

        val blocked = eventStore.read(sessionId).mapNotNull { it.payload as? PreemptRedirectBlockedEvent }
        assertEquals(1, blocked.size)
        assertTrue(blocked.single().reason.contains("unsatisfied needs"), blocked.single().reason)
        // Normal path was taken: the first hop out of A went to B, not C.
        val firstFromA = transitions(sessionId).first { it.from == StageId("A") }
        assertEquals(StageId("B"), firstFromA.to)
        assertNull(firstFromA.redirectId)
    }

    @Test
    fun `redirect to an unknown stage is blocked`(): Unit = runBlocking {
        val sessionId = SessionId("reroute-block-unknown")
        append(
            sessionId,
            PreemptRedirectEvent(
                redirectId = "r3",
                sessionId = sessionId,
                fromStageId = StageId("A"),
                toStageId = StageId("ZZZ"),
                timestampMs = 1L,
            ),
        )

        orchestrator().run(sessionId, linearGraph(), config)

        val blocked = eventStore.read(sessionId).mapNotNull { it.payload as? PreemptRedirectBlockedEvent }
        assertEquals(1, blocked.size)
        assertTrue(blocked.single().reason.contains("unknown stage"), blocked.single().reason)
    }

    @Test
    fun `an already-consumed redirect is ignored`(): Unit = runBlocking {
        val sessionId = SessionId("reroute-consumed")
        // Redirect whose id already appears on a prior TransitionExecuted → treated as applied.
        append(
            sessionId,
            PreemptRedirectEvent(
                redirectId = "r4",
                sessionId = sessionId,
                fromStageId = StageId("A"),
                toStageId = StageId("C"),
                timestampMs = 1L,
            ),
        )
        append(
            sessionId,
            TransitionExecutedEvent(sessionId, StageId("A"), StageId("C"), TransitionId("redirect:r4"), redirectId = "r4"),
        )

        orchestrator().run(sessionId, linearGraph(), config)

        // No block recorded, and the live run resolves A→B normally (the stale redirect is inert).
        assertTrue(eventStore.read(sessionId).none { it.payload is PreemptRedirectBlockedEvent })
        val firstFromA = transitions(sessionId).last { it.from == StageId("A") }
        assertEquals(StageId("B"), firstFromA.to)
    }
}
