import com.correx.core.approvals.ApprovalProjector
import com.correx.core.approvals.DefaultApprovalReducer
import com.correx.core.approvals.DefaultApprovalRepository
import com.correx.core.artifacts.DefaultArtifactReducer
import com.correx.core.artifactstore.ArtifactStore
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.NewEvent
import com.correx.core.events.events.SessionWorkspaceBoundEvent
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.orchestration.OrchestrationState
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.inference.InferenceProvider
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.kernel.execution.ReplayStrategy
import com.correx.core.kernel.orchestration.OrchestrationConfig
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.orchestration.ReplayOrchestrator
import com.correx.core.kernel.orchestration.WorkspaceContext
import com.correx.core.kernel.orchestration.WorkspaceToolRegistryProvider
import com.correx.core.kernel.orchestration.WorkspaceTools
import com.correx.core.risk.NoOpRiskAssessor
import com.correx.core.sessions.DefaultSessionReducer
import com.correx.core.sessions.DefaultSessionRepository
import com.correx.core.sessions.SessionProjector
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.sessions.projections.replay.EventReplayer
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.registry.ToolRegistry
import com.correx.core.transitions.evaluation.EvaluationContext
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.TransitionCondition
import com.correx.core.transitions.graph.TransitionEdge
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.core.transitions.resolution.TransitionDecision
import com.correx.core.transitions.resolution.TransitionResolver
import com.correx.core.utils.TypeId
import com.correx.core.validation.pipeline.ValidationPipeline
import com.correx.infrastructure.persistence.InMemoryEventStore
import com.correx.infrastructure.persistence.artifact.LiveArtifactRepository
import com.correx.testing.fixtures.InferenceFixtures
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test
import java.nio.file.Path
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicReference

class ReplayWorkspaceDerivationTest {

    private val sessionId = SessionId("s-replay-ws")
    private val stageId = StageId("stage-1")

    /** Spy: records the WorkspaceContext passed to forWorkspace(). */
    private val capturedWorkspace = AtomicReference<WorkspaceContext?>(null)
    private val spyProvider = WorkspaceToolRegistryProvider { ws ->
        capturedWorkspace.set(ws)
        WorkspaceTools(
            registry = object : ToolRegistry {
                override fun resolve(name: String): Tool? = null
                override fun all(): List<Tool> = emptyList()
            },
            executor = object : ToolExecutor {
                override suspend fun execute(request: ToolRequest): ToolResult =
                    ToolResult.Success(invocationId = request.invocationId, output = "")
            },
        )
    }

    @Test
    fun `replay derives workspace from SessionWorkspaceBoundEvent`(): Unit = runBlocking {
        val store = InMemoryEventStore()
        val recordedRoot = "/recorded/workspace"
        seedSessionWorkspaceBoundEvent(store, workspaceRoot = recordedRoot)

        val orchestrator = buildOrchestrator(store)
        val config = buildConfig(defaultWorkspaceRoot = "/default/workspace")

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

        assertEquals(recordedRoot, capturedWorkspace.get()?.workspaceRoot?.toString())
    }

    @Test
    fun `replay falls back to config workspace when no SessionWorkspaceBoundEvent`(): Unit = runBlocking {
        val store = InMemoryEventStore()
        val defaultRoot = "/default/workspace"

        val orchestrator = buildOrchestrator(store)
        val config = buildConfig(defaultWorkspaceRoot = defaultRoot)

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

        assertEquals(defaultRoot, capturedWorkspace.get()?.workspaceRoot?.toString())
    }

    // ---- helpers ----

    private suspend fun seedSessionWorkspaceBoundEvent(store: InMemoryEventStore, workspaceRoot: String) {
        store.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId("e-ws"),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = SessionWorkspaceBoundEvent(
                    sessionId = sessionId,
                    workspaceRoot = workspaceRoot,
                    allowedPaths = listOf(workspaceRoot),
                ),
            )
        )
    }

    private fun buildConfig(defaultWorkspaceRoot: String) = OrchestrationConfig(
        replayStrategy = ReplayStrategy.SkipInference,
        workspace = WorkspaceContext(
            workspaceRoot = Path.of(defaultWorkspaceRoot),
            workingDir = Path.of(defaultWorkspaceRoot),
            allowedPaths = setOf(Path.of(defaultWorkspaceRoot)),
        ),
    )

    private fun buildGraph(): WorkflowGraph {
        val terminal = StageId("terminal")
        return WorkflowGraph(
            id = "test-wf",
            start = stageId,
            stages = mapOf(
                stageId to StageConfig(produces = emptyList(), allowedTools = emptySet()),
                terminal to StageConfig(produces = emptyList(), allowedTools = emptySet()),
            ),
            transitions = setOf(
                TransitionEdge(
                    id = com.correx.core.events.types.TransitionId("t1"),
                    from = stageId,
                    to = terminal,
                    condition = TransitionCondition { _ -> true },
                )
            ),
        )
    }

    private fun buildOrchestrator(store: InMemoryEventStore): ReplayOrchestrator {
        val sessionReplayer = DefaultEventReplayer(store, SessionProjector(DefaultSessionReducer()))
        val sessionRepository = DefaultSessionRepository(sessionReplayer)

        val artifactStore = object : ArtifactStore {
            override suspend fun put(bytes: ByteArray): ArtifactId =
                TypeId("00".repeat(32))
            override suspend fun get(id: ArtifactId): ByteArray? = null
            override suspend fun flushBefore(commit: suspend () -> Unit) { commit() }
        }

        val inferenceReplayer = object : EventReplayer<InferenceState> {
            override fun rebuild(sessionId: SessionId) = InferenceState()
        }
        val orchReplayer = object : EventReplayer<OrchestrationState> {
            override fun rebuild(sessionId: SessionId) = OrchestrationState()
        }

        val repositories = OrchestratorRepositories(
            eventStore = store,
            inferenceRepository = InferenceRepository(inferenceReplayer),
            orchestrationRepository = OrchestrationRepository(orchReplayer),
            sessionRepository = sessionRepository,
            artifactRepository = LiveArtifactRepository(store, DefaultArtifactReducer()),
            approvalRepository = DefaultApprovalRepository(
                DefaultEventReplayer(store, ApprovalProjector(DefaultApprovalReducer()))
            ),
        )

        val engines = OrchestratorEngines(
            transitionResolver = object : TransitionResolver {
                override fun resolve(graph: WorkflowGraph, context: EvaluationContext): TransitionDecision =
                    TransitionDecision.Move(
                        to = StageId("terminal"),
                        transitionId = com.correx.core.events.types.TransitionId("t1"),
                    )
            },
            contextPackBuilder = object : ContextPackBuilder {
                override suspend fun build(
                    id: ContextPackId,
                    sessionId: SessionId,
                    stageId: StageId,
                    entries: List<ContextEntry>,
                    budget: TokenBudget,
                ): ContextPack = InferenceFixtures.contextPack()
            },
            inferenceRouter = object : InferenceRouter {
                override suspend fun route(
                    stageId: StageId,
                    requiredCapabilities: Set<ModelCapability>,
                ): InferenceProvider = throw UnsupportedOperationException("should use replay provider")
            },
            validationPipeline = ValidationPipeline(validators = emptyList()),
            approvalEngine = com.correx.core.approvals.domain.NoOpApprovalEngine(),
            riskAssessor = NoOpRiskAssessor(),
            workspaceToolRegistryProvider = spyProvider,
        )

        return ReplayOrchestrator(
            repositories = repositories,
            engines = engines,
            cancellations = ConcurrentHashMap(),
            strategy = ReplayStrategy.SkipInference,
            artifactStore = artifactStore,
        )
    }
}
