package com.correx.infrastructure.workflow

import com.correx.core.events.events.LspDiagnostic
import com.correx.core.kernel.orchestration.LspDiagnosticsRequest
import com.correx.core.kernel.orchestration.LspDiagnosticsResult
import com.correx.core.kernel.orchestration.LspDiagnosticsRunner
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.eclipse.lsp4j.ClientCapabilities
import org.eclipse.lsp4j.Diagnostic
import org.eclipse.lsp4j.DidOpenTextDocumentParams
import org.eclipse.lsp4j.InitializeParams
import org.eclipse.lsp4j.InitializedParams
import org.eclipse.lsp4j.MessageActionItem
import org.eclipse.lsp4j.MessageParams
import org.eclipse.lsp4j.PublishDiagnosticsCapabilities
import org.eclipse.lsp4j.PublishDiagnosticsParams
import org.eclipse.lsp4j.ShowMessageRequestParams
import org.eclipse.lsp4j.TextDocumentClientCapabilities
import org.eclipse.lsp4j.TextDocumentItem
import org.eclipse.lsp4j.launch.LSPLauncher
import org.eclipse.lsp4j.services.LanguageClient
import java.nio.file.Files
import java.util.concurrent.CompletableFuture
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.TimeUnit
import kotlin.coroutines.cancellation.CancellationException

data class LspServerSpec(val id: String, val command: List<String>, val languageId: String)

class LspServerRegistry(
    private val byExtension: Map<String, LspServerSpec> = defaults,
) {
    fun serverFor(path: String): LspServerSpec? = byExtension[path.substringAfterLast('.', "").lowercase()]

    companion object {
        private val typescript = LspServerSpec(
            "tsserver",
            listOf("typescript-language-server", "--stdio"),
            "typescriptreact",
        )
        private val rust = LspServerSpec("rust-analyzer", listOf("rust-analyzer"), "rust")
        private val go = LspServerSpec("gopls", listOf("gopls"), "go")
        private val clangd = LspServerSpec("clangd", listOf("clangd"), "cpp")
        val defaults = mapOf(
            "ts" to typescript.copy(languageId = "typescript"),
            "tsx" to typescript,
            "js" to typescript.copy(languageId = "javascript"),
            "jsx" to typescript.copy(languageId = "javascriptreact"),
            "rs" to rust,
            "go" to go,
            "c" to clangd.copy(languageId = "c"),
            "cc" to clangd,
            "cpp" to clangd,
            "h" to clangd.copy(languageId = "c"),
            "hpp" to clangd,
        )
    }
}

/**
 * LSP push-diagnostics adapter: `didOpen` each file, then collect the `publishDiagnostics`
 * the server pushes back. tsserver only supports push (rejects the 3.17 pull method);
 * rust-analyzer/gopls/clangd push on open too. Missing servers degrade to the static floor.
 */
class Lsp4jDiagnosticsRunner(
    private val registry: LspServerRegistry = LspServerRegistry(),
    private val timeoutSeconds: Long = 30,
) : LspDiagnosticsRunner {
    override suspend fun pull(request: LspDiagnosticsRequest): LspDiagnosticsResult = withContext(Dispatchers.IO) {
        val groups = request.paths
            .mapNotNull { path -> registry.serverFor(path)?.let { it to path } }
            .groupBy({ it.first }, { it.second })
        if (groups.isEmpty()) {
            return@withContext LspDiagnosticsResult(skippedReason = "no language server registered for touched files")
        }
        val diagnostics = mutableListOf<LspDiagnostic>()
        val skipped = mutableListOf<String>()
        groups.forEach { (server, paths) ->
            runCatching { diagnostics += pullFromServer(request, server, paths) }
                .onFailure {
                    if (it is CancellationException) throw it
                    skipped += "${server.id}: ${it.message ?: it::class.simpleName}"
                }
        }
        LspDiagnosticsResult(
            server = groups.keys.joinToString(",") { it.id },
            diagnostics = diagnostics,
            skippedReason = skipped.takeIf { it.isNotEmpty() }?.joinToString("; "),
        )
    }

    private fun pullFromServer(
        request: LspDiagnosticsRequest,
        spec: LspServerSpec,
        paths: List<String>,
    ): List<LspDiagnostic> {
        val process = ProcessBuilder(spec.command).directory(request.workspaceRoot.toFile()).start()
        val stderrDrain = Thread.ofVirtual().start { process.errorStream.bufferedReader().use { it.readText() } }
        val client = CollectingLanguageClient()
        val launcher = LSPLauncher.createClientLauncher(client, process.inputStream, process.outputStream)
        val listening = launcher.startListening()
        val server = launcher.remoteProxy
        return try {
            server.initialize(
                InitializeParams().apply {
                    rootUri = request.workspaceRoot.toUri().toString()
                    capabilities = ClientCapabilities().apply {
                        textDocument = TextDocumentClientCapabilities().apply {
                            publishDiagnostics = PublishDiagnosticsCapabilities()
                        }
                    }
                },
            ).get(timeoutSeconds, TimeUnit.SECONDS)
            server.initialized(InitializedParams())
            val uriToPath = paths.associateBy { path ->
                request.workspaceRoot.resolve(path).normalize().toUri().toString()
            }
            uriToPath.forEach { (uri, path) ->
                val file = request.workspaceRoot.resolve(path).normalize()
                server.textDocumentService.didOpen(
                    DidOpenTextDocumentParams(TextDocumentItem(uri, spec.languageId, 1, Files.readString(file))),
                )
            }
            awaitDiagnostics(client, uriToPath.keys)
            uriToPath.flatMap { (uri, path) ->
                client.latestFor(uri).map { diagnostic ->
                    LspDiagnostic(
                        path = path,
                        line = diagnostic.range.start.line,
                        character = diagnostic.range.start.character,
                        severity = diagnostic.severity?.name?.lowercase() ?: "error",
                        code = diagnostic.code?.let { if (it.isLeft) it.left else it.right.toString() },
                        message = diagnostic.message,
                        tags = diagnostic.tags.orEmpty().map { it.name.lowercase() },
                    )
                }
            }
        } finally {
            runCatching { server.shutdown().get(timeoutSeconds, TimeUnit.SECONDS) }
            server.exit()
            listening.cancel(true)
            process.destroy()
            stderrDrain.join(TimeUnit.SECONDS.toMillis(1))
        }
    }

    /**
     * Block until every opened URI has received at least one push (or the timeout elapses), then
     * wait for a quiet period anchored to the last server publication. Servers such as tsserver
     * can publish an empty report while loading the project and real diagnostics afterward; a
     * timer started immediately after didOpen reads the workspace before the server is ready.
     */
    private fun awaitDiagnostics(client: CollectingLanguageClient, uris: Set<String>) {
        val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(timeoutSeconds)
        client.awaitAll(uris, deadline)
        client.awaitQuiet(deadline, TimeUnit.MILLISECONDS.toNanos(SETTLE_MILLIS))
    }

    private class CollectingLanguageClient : LanguageClient {
        private val byUri = ConcurrentHashMap<String, List<Diagnostic>>()
        private val updates = Object()
        @Volatile private var lastUpdateNanos = 0L
        fun latestFor(uri: String): List<Diagnostic> = byUri[uri].orEmpty()
        fun hasAll(uris: Set<String>): Boolean = byUri.keys.containsAll(uris)

        fun awaitAll(uris: Set<String>, deadline: Long) {
            synchronized(updates) {
                while (!hasAll(uris)) {
                    val remaining = deadline - System.nanoTime()
                    if (remaining <= 0) return
                    updates.wait(TimeUnit.NANOSECONDS.toMillis(remaining).coerceAtLeast(1))
                }
            }
        }

        fun awaitQuiet(deadline: Long, quietNanos: Long) {
            synchronized(updates) {
                while (true) {
                    val remaining = deadline - System.nanoTime()
                    if (remaining <= 0) return
                    val sinceUpdate = System.nanoTime() - lastUpdateNanos
                    if (lastUpdateNanos != 0L && sinceUpdate >= quietNanos) return
                    val waitNanos = minOf(remaining, quietNanos - sinceUpdate)
                    updates.wait(TimeUnit.NANOSECONDS.toMillis(waitNanos).coerceAtLeast(1))
                }
            }
        }

        override fun publishDiagnostics(diagnostics: PublishDiagnosticsParams) {
            byUri[diagnostics.uri] = diagnostics.diagnostics.orEmpty()
            synchronized(updates) {
                lastUpdateNanos = System.nanoTime()
                updates.notifyAll()
            }
        }
        override fun telemetryEvent(`object`: Any?) = Unit
        override fun showMessage(messageParams: MessageParams) = Unit
        override fun showMessageRequest(requestParams: ShowMessageRequestParams): CompletableFuture<MessageActionItem> =
            CompletableFuture.completedFuture(null)
        override fun logMessage(message: MessageParams) = Unit
    }

    private companion object {
        const val SETTLE_MILLIS = 750L
    }
}
