package com.correx.apps.server.ws

import com.correx.apps.server.ServerModule
import com.correx.apps.server.configureServer
import com.correx.apps.server.protocol.ClientMessage
import com.correx.apps.server.registry.ProviderRegistry
import com.correx.apps.server.registry.WorkflowRegistry
import com.correx.apps.server.registry.WorkflowSummary
import com.correx.apps.server.undo.SessionUndoService
import com.correx.apps.server.workspace.WorkspaceResolver
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.SessionWorkspaceBoundEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.StageId
import com.correx.core.inference.DefaultInferenceRouter
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ProviderHealth
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.orchestration.WorkspaceContext
import com.correx.core.kernel.retry.DefaultRetryCoordinator
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.transitions.graph.StageConfig
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.infrastructure.InfrastructureModule
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
import com.correx.infrastructure.inference.commons.UnavailableProbe
import com.correx.infrastructure.persistence.InMemoryEventStore
import io.ktor.client.plugins.websocket.WebSockets
import io.ktor.client.plugins.websocket.webSocket
import io.ktor.server.testing.testApplication
import io.ktor.websocket.Frame
import kotlinx.coroutines.delay
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.io.TempDir
import java.nio.file.Path
import java.util.concurrent.ConcurrentLinkedQueue

/**
 * Focused integration tests for the Axis-2 Phase-B workspace handshake.
 *
 * Protocol sequence under test:
 *   Hello(workingDir) → StartSession(workflowId) → SessionWorkspaceBoundEvent emitted
 *
 * Absent-Hello fallback:
 *   StartSession(workflowId) → no crash, no SessionWorkspaceBoundEvent (resolver not called)
 *
 * Note: the orchestrator is not driven to completion; we only assert on the event log
 * after the handshake is processed and before the orchestrator runs far enough to interfere.
 */
class WorkspaceHandshakeTest {

    private val protocolJson = Json { classDiscriminator = "type"; ignoreUnknownKeys = true }
    private fun encode(msg: ClientMessage): String = protocolJson.encodeToString(msg)

    // ---- infrastructure helpers ------------------------------------------------

    private val noopArtifactStore: ArtifactStore = object : ArtifactStore {
        override suspend fun put(bytes: ByteArray): ArtifactId = ArtifactId("noop")
        override suspend fun get(id: ArtifactId): ByteArray? = null
        override suspend fun flushBefore(commit: suspend () -> Unit) = commit()
    }

    private val noopProviderRegistry: ProviderRegistry = object : ProviderRegistry {
        override fun listAll() = emptyList<InferenceProvider>()
        override suspend fun healthCheckAll() = emptyMap<ProviderId, ProviderHealth>()
    }

    /**
     * An event store that delegates to [InMemoryEventStore] but also copies each appended
     * event into [appended] for inspection.
     */
    private fun capturingEventStore(): Pair<EventStore, ConcurrentLinkedQueue<NewEvent>> {
        val delegate = InMemoryEventStore()
        val captured = ConcurrentLinkedQueue<NewEvent>()
        val store = object : EventStore by delegate {
            override suspend fun append(event: NewEvent): StoredEvent {
                captured.add(event)
                return delegate.append(event)
            }
        }
        return store to captured
    }

    private fun buildWorkflowRegistry(graph: WorkflowGraph): WorkflowRegistry = object : WorkflowRegistry {
        override fun listAll() = listOf(WorkflowSummary(graph.id, description = ""))
        override fun find(workflowId: String) = if (workflowId == graph.id) graph else null
    }

    private fun minimalGraph(): WorkflowGraph = WorkflowGraph(
        id = "handshake-test-workflow",
        stages = mapOf(StageId("s1") to StageConfig(allowedTools = emptySet())),
        transitions = emptySet(),
        start = StageId("s1"),
    )

    private fun buildModule(
        eventStore: EventStore,
        workflowRegistry: WorkflowRegistry,
        workspaceResolver: WorkspaceResolver?,
        tempDir: Path,
    ): ServerModule {
        val provider = InfrastructureModule.createLlamaCppProvider(
            modelId = "test-model",
            modelPath = "/dev/null",
            baseUrl = "http://127.0.0.1:1",
        )
        val infraRegistry = InfrastructureModule.createProviderRegistry(listOf(provider))
        val inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy())

        val toolConfig = com.correx.infrastructure.tools.ToolConfig(
            shell = com.correx.infrastructure.tools.ShellConfig(
                enabled = false, allowedExecutables = emptySet(), workingDir = tempDir,
            ),
            fileRead = com.correx.infrastructure.tools.FileReadConfig(
                enabled = false, allowedPaths = setOf(tempDir),
            ),
            fileWrite = com.correx.infrastructure.tools.FileWriteConfig(
                enabled = false, allowedPaths = setOf(tempDir), workingDir = tempDir,
            ),
            fileEdit = com.correx.infrastructure.tools.FileEditConfig(
                enabled = false, allowedPaths = setOf(tempDir), workingDir = tempDir,
            ),
        )
        val toolRegistry = InfrastructureModule.createToolRegistry(toolConfig)
        val eventDispatcher = com.correx.core.events.EventDispatcher(eventStore)
        val toolExecutor = InfrastructureModule.createToolExecutor(
            registry = toolRegistry,
            eventDispatcher = eventDispatcher,
            workDir = tempDir,
            artifactStore = null,
        )

        val engines = OrchestratorEngines(
            transitionResolver = com.correx.core.transitions.resolution.DefaultTransitionResolver {
                    condition, ctx -> condition.evaluate(ctx)
            },
            contextPackBuilder = com.correx.core.context.builder.DefaultContextPackBuilder(
                com.correx.core.context.compression.DefaultContextCompressor(),
            ),
            inferenceRouter = inferenceRouter,
            validationPipeline = com.correx.core.validation.pipeline.ValidationPipeline(validators = emptyList()),
            approvalEngine = com.correx.core.approvals.domain.DefaultApprovalEngine(),
            riskAssessor = com.correx.core.risk.DefaultRiskAssessor(),
            toolRegistry = toolRegistry,
            toolExecutor = toolExecutor,
        )

        val repositories = OrchestratorRepositories(
            eventStore = eventStore,
            inferenceRepository = com.correx.core.inference.InferenceRepository(
                DefaultEventReplayer(eventStore, com.correx.core.inference.InferenceProjector()),
            ),
            orchestrationRepository = OrchestrationRepository(
                DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())),
            ),
            sessionRepository = DefaultSessionRepository(
                DefaultEventReplayer(eventStore, SessionProjector(DefaultSessionReducer())),
            ),
            artifactRepository = InfrastructureModule.createArtifactRepository(eventStore),
            approvalRepository = InfrastructureModule.createApprovalRepository(eventStore),
        )

        val orchestrator = DefaultSessionOrchestrator(
            repositories = repositories,
            engines = engines,
            retryCoordinator = DefaultRetryCoordinator(eventStore),
            artifactStore = noopArtifactStore,
            tokenizer = provider.tokenizer,
            decisionJournalRepository = InfrastructureModule.createDecisionJournalRepository(eventStore),
        )

        val routerFacade = InfrastructureModule.createTalkieFacade(
            eventStore = eventStore,
            inferenceRouter = inferenceRouter,
            config = com.correx.core.talkie.model.TalkieConfig(),
            tokenizer = provider.tokenizer,
        )

        val sessionUndoService = SessionUndoService(
            eventStore = eventStore,
            artifactStore = noopArtifactStore,
            bootRoots = setOf(tempDir),
        )

        return ServerModule(
            orchestrator = orchestrator,
            eventStore = eventStore,
            artifactStore = noopArtifactStore,
            sessionRepository = repositories.sessionRepository,
            workflowRegistry = workflowRegistry,
            providerRegistry = noopProviderRegistry,
            defaultOrchestrationConfig = OrchestrationConfig(sandboxRoot = tempDir),
            routerFacade = routerFacade,
            orchestrationRepository = repositories.orchestrationRepository,
            approvalRepository = repositories.approvalRepository,
            toolRegistry = toolRegistry,
            sessionUndoService = sessionUndoService,
            resourceProbe = UnavailableProbe,
            workspaceResolver = workspaceResolver,
        )
    }

    // ---- tests -----------------------------------------------------------------

    @Test
    fun `Hello followed by StartSession emits SessionWorkspaceBoundEvent with resolved workspace`(
        @TempDir tempDir: Path,
    ) {
        val (eventStore, captured) = capturingEventStore()
        val bootWorkspace = WorkspaceContext(
            workspaceRoot = tempDir,
            workingDir = tempDir,
            allowedPaths = setOf(tempDir),
        )
        val resolver = WorkspaceResolver(
            bootDefault = bootWorkspace,
            privilegedLocations = emptyList(),
            allowedWorkspaceRoots = emptyList(),
        )
        val module = buildModule(eventStore, buildWorkflowRegistry(minimalGraph()), resolver, tempDir)
        module.start()

        testApplication {
            application { configureServer(module) }
            val client = createClient { install(WebSockets) }

            client.webSocket("/stream") {
                // Drain initial snapshot (WorkflowList, provider health, ResourceStatus)
                delay(100)

                // Phase 1: send Hello with the tempDir as workingDir
                send(Frame.Text(encode(ClientMessage.Hello(workingDir = tempDir.toString()))))
                delay(50)

                // Phase 2: start a session
                val startMsg = ClientMessage.StartSession(workflowId = "handshake-test-workflow", config = null)
                send(Frame.Text(encode(startMsg)))

                // Give the server time to process the event emission (before orchestrator runs)
                delay(200)
            }
        }

        val bound = captured.firstOrNull { it.payload is SessionWorkspaceBoundEvent }
        assertNotNull(bound, "SessionWorkspaceBoundEvent must be emitted after Hello + StartSession")

        val payload = bound!!.payload as SessionWorkspaceBoundEvent
        assert(payload.workspaceRoot == tempDir.toRealPath().toString()) {
            "workspaceRoot must equal resolved tempDir, got: ${payload.workspaceRoot}"
        }
    }

    @Test
    fun `Hello after StartSession is ignored - first session workspace is unaffected`(
        @TempDir tempDir: Path,
    ) {
        // Sequence: Hello(dir1) → StartSession → Hello(dir2) [late, must be ignored]
        // The SessionWorkspaceBoundEvent must record dir1's resolution, not dir2.
        val (eventStore, captured) = capturingEventStore()
        val bootWorkspace = WorkspaceContext(
            workspaceRoot = tempDir,
            workingDir = tempDir,
            allowedPaths = setOf(tempDir),
        )
        val resolver = WorkspaceResolver(
            bootDefault = bootWorkspace,
            privilegedLocations = emptyList(),
            allowedWorkspaceRoots = emptyList(),
        )
        val module = buildModule(eventStore, buildWorkflowRegistry(minimalGraph()), resolver, tempDir)
        module.start()

        testApplication {
            application { configureServer(module) }
            val client = createClient { install(WebSockets) }

            client.webSocket("/stream") {
                delay(100)

                // Hello with the real tempDir
                send(Frame.Text(encode(ClientMessage.Hello(workingDir = tempDir.toString()))))
                delay(50)

                // StartSession — this closes the Hello window
                val startMsg = ClientMessage.StartSession(workflowId = "handshake-test-workflow", config = null)
                send(Frame.Text(encode(startMsg)))
                delay(100)

                // Late Hello with a different path — must be ignored
                send(Frame.Text(encode(ClientMessage.Hello(workingDir = "/tmp/late-hello-should-be-ignored"))))
                delay(200)
            }
        }

        // Exactly one SessionWorkspaceBoundEvent, recording tempDir (not the late Hello path)
        val boundEvents = captured.filter { it.payload is SessionWorkspaceBoundEvent }
        assert(boundEvents.size == 1) { "Expected exactly 1 SessionWorkspaceBoundEvent, got ${boundEvents.size}" }
        val payload = boundEvents.first().payload as SessionWorkspaceBoundEvent
        assert(payload.workspaceRoot == tempDir.toRealPath().toString()) {
            "workspaceRoot must equal first Hello's resolved dir, got: ${payload.workspaceRoot}"
        }
    }

    @Test
    fun `absent Hello causes graceful fallback - no crash and no SessionWorkspaceBoundEvent when resolver is null`(
        @TempDir tempDir: Path,
    ) {
        // workspaceResolver = null simulates the legacy no-resolver path (Phase A default).
        // No Hello is sent. The server must not crash and must not emit SessionWorkspaceBoundEvent.
        val (eventStore, captured) = capturingEventStore()
        val module = buildModule(eventStore, buildWorkflowRegistry(minimalGraph()), null, tempDir)
        module.start()

        testApplication {
            application { configureServer(module) }
            val client = createClient { install(WebSockets) }

            client.webSocket("/stream") {
                delay(100)

                // No Hello — send StartSession directly
                val startMsg = ClientMessage.StartSession(workflowId = "handshake-test-workflow", config = null)
                send(Frame.Text(encode(startMsg)))
                delay(200)
            }
        }

        val bound = captured.firstOrNull { it.payload is SessionWorkspaceBoundEvent }
        assertNull(bound, "No SessionWorkspaceBoundEvent must be emitted when workspaceResolver is null")
    }
}
