CodeMapper/Sources/CodeMapper/LSPClient.swift

import Foundation

enum LSPError: Error, CustomStringConvertible {
    case connectionClosed
    case invalidHeader
    case invalidJSON
    case requestFailed(String)

    var description: String {
        switch self {
        case .connectionClosed: return "LSP connection closed"
        case .invalidHeader: return "Invalid LSP header"
        case .invalidJSON: return "Invalid LSP JSON"
        case .requestFailed(let msg): return "LSP request failed: \(msg)"
        }
    }
}

enum LSPServerMode {
    /// sourcekit-lsp — used for Swift packages. clangd flags are forwarded
    /// via `-Xclangd`. Index-store discovery is attempted for faster call graphs.
    case sourcekitLsp
    /// clangd directly — used for pure C/C++ projects (e.g. CMake). Flags are
    /// passed directly; no index-store discovery.
    case clangd(compileCommandsDir: String)
}

class LSPClient {
    private let process: Process
    private let stdinPipe: Pipe
    private let stdoutPipe: Pipe
    private var nextId = 1
    private var isShutdown = false
    let projectRoot: String
    private let serverMode: LSPServerMode

    /// Diagnostics pushed by the server via `textDocument/publishDiagnostics`
    /// notifications, keyed by document URI. The server computes these
    /// asynchronously after `didOpen`/`didChange`; we capture them
    /// opportunistically while reading responses (see `sendRequest`) since
    /// there's no synchronous "give me diagnostics now" LSP request. In
    /// practice clangd/sourcekit-lsp publish them promptly enough that they
    /// arrive interleaved with (or before) the next request's response.
    private var diagnosticsByURI: [String: [[String: Any]]] = [:]

    init(lspPath: String, projectRoot: String, compileCommandsDir: String? = nil, serverMode: LSPServerMode = .sourcekitLsp) throws {
        self.projectRoot = projectRoot
        self.serverMode = serverMode

        process = Process()
        process.executableURL = URL(fileURLWithPath: lspPath)

        switch serverMode {
        case .sourcekitLsp:
            if let compileCommandsDir {
                // Forwarded through sourcekit-lsp to its internal clangd so
                // C-family files are parsed with the correct driver language
                // (see CompilationDatabaseWriter). We deliberately don't force
                // `--default-workspace-type compilationDatabase` — that bypasses
                // SwiftPM build-system integration entirely, which would degrade
                // Swift-side indexing/call-hierarchy in mixed Swift+C-family
                // packages (the primary use case) for no diagnostic benefit:
                // sourcekit-lsp doesn't forward clangd's publishDiagnostics for
                // C-family files in either workspace mode (verified empirically —
                // not even injected syntax errors produce a notification).
                process.arguments = ["-Xclangd", "-compile-commands-dir=\(compileCommandsDir)"]
            }
        case .clangd(let compileCommandsDir):
            process.arguments = ["--compile-commands-dir=\(compileCommandsDir)"]
        }

        stdinPipe = Pipe()
        stdoutPipe = Pipe()

        process.standardInput = stdinPipe
        process.standardOutput = stdoutPipe
        // Log LSP stderr to /tmp for debugging
        let logPath = "/tmp/sourcekit-lsp.log"
        FileManager.default.createFile(atPath: logPath, contents: nil)
        process.standardError = FileHandle(forWritingAtPath: logPath)

        try process.run()
    }

    deinit {
        if !isShutdown { process.terminate() }
    }

    func initialize() throws {
        var initOptions: [String: Any] = [:]

        // Index-store discovery only applies to sourcekit-lsp (Swift packages).
        // clangd builds its own index from compile_commands.json.
        if case .sourcekitLsp = serverMode {
            if let isp = findIndexStorePath() {
                initOptions["indexStorePath"] = isp
                initOptions["indexDatabasePath"] = "/tmp/codemapper-lsp-index"
            }
        }

        let params: [String: Any] = [
            "processId": ProcessInfo.processInfo.processIdentifier,
            "rootUri": "file://\(projectRoot)",
            "capabilities": [
                "textDocument": [
                    "callHierarchy": ["dynamicRegistration": false],
                    "documentSymbol": ["hierarchicalDocumentSymbolSupport": true]
                ]
            ],
            "initializationOptions": initOptions
        ]
        _ = try sendRequest("initialize", params: params)
        try sendNotification("initialized", params: [:])

        if case .sourcekitLsp = serverMode, !initOptions.isEmpty {
            // Brief pause for sourcekit-lsp to load the index store
            Thread.sleep(forTimeInterval: 2.0)
        }
    }

    private func findIndexStorePath() -> String? {
        let fm = FileManager.default
        let candidates = [
            projectRoot + "/.build/arm64-apple-macosx/debug/index/store",
            projectRoot + "/.build/debug/index/store",
            projectRoot + "/.build/index-build/arm64-apple-macosx/debug/index/store",
        ]
        return candidates.first { fm.fileExists(atPath: $0) }
    }

    func openFile(uri: String, content: String, languageId: String = "swift") throws {
        try sendNotification("textDocument/didOpen", params: [
            "textDocument": [
                "uri": uri,
                "languageId": languageId,
                "version": 1,
                "text": content
            ] as [String: Any]
        ])
    }

    func closeFile(uri: String) throws {
        try sendNotification("textDocument/didClose", params: [
            "textDocument": ["uri": uri]
        ])
    }

    func documentSymbol(uri: String) throws -> [[String: Any]] {
        let result = try sendRequest("textDocument/documentSymbol", params: [
            "textDocument": ["uri": uri]
        ])
        return result as? [[String: Any]] ?? []
    }

    func callHierarchyPrepare(uri: String, line: Int, character: Int) throws -> [[String: Any]] {
        let result = try sendRequest("textDocument/prepareCallHierarchy", params: [
            "textDocument": ["uri": uri],
            "position": ["line": line, "character": character]
        ])
        return result as? [[String: Any]] ?? []
    }

    func callHierarchyOutgoing(item: [String: Any]) throws -> [[String: Any]] {
        let result = try sendRequest("callHierarchy/outgoingCalls", params: [
            "item": item
        ])
        return result as? [[String: Any]] ?? []
    }

    func shutdownGracefully() throws {
        guard !isShutdown else { return }
        isShutdown = true
        _ = try? sendRequest("shutdown", params: [:] as [String: Any])
        try? sendNotification("exit", params: [:])
        process.terminate()
    }

    // MARK: - JSON-RPC transport

    @discardableResult
    func sendRequest(_ method: String, params: Any) throws -> Any? {
        let id = nextId
        nextId += 1

        let msg: [String: Any] = [
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params
        ]
        try writeMessage(msg)

        while true {
            let response = try readMessage()
            // Check if this is our response (has matching id)
            if let respId = response["id"] as? Int, respId == id {
                if let error = response["error"] as? [String: Any] {
                    let msg = error["message"] as? String ?? "unknown error"
                    throw LSPError.requestFailed(msg)
                }
                return response["result"]
            }
            recordDiagnosticsIfPresent(response)
            // Discard other notifications / non-matching responses and keep reading
        }
    }

    private func recordDiagnosticsIfPresent(_ message: [String: Any]) {
        guard message["method"] as? String == "textDocument/publishDiagnostics",
              let params = message["params"] as? [String: Any],
              let uri = params["uri"] as? String,
              let diagnostics = params["diagnostics"] as? [[String: Any]]
        else { return }
        diagnosticsByURI[uri] = diagnostics
    }

    /// Returns and clears any diagnostics captured for `uri` since the last call.
    func takeDiagnostics(uri: String) -> [[String: Any]] {
        diagnosticsByURI.removeValue(forKey: uri) ?? []
    }

    private func sendNotification(_ method: String, params: Any) throws {
        let msg: [String: Any] = [
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        ]
        try writeMessage(msg)
    }

    private func writeMessage(_ msg: [String: Any]) throws {
        let body = try JSONSerialization.data(withJSONObject: msg)
        let header = "Content-Length: \(body.count)\r\n\r\n"
        var data = header.data(using: .utf8)!
        data.append(body)
        stdinPipe.fileHandleForWriting.write(data)
    }

    private func readMessage() throws -> [String: Any] {
        let contentLength = try readContentLength()
        let body = stdoutPipe.fileHandleForReading.readData(ofLength: contentLength)
        guard body.count == contentLength else { throw LSPError.connectionClosed }
        guard let obj = try JSONSerialization.jsonObject(with: body) as? [String: Any] else {
            throw LSPError.invalidJSON
        }
        return obj
    }

    private func readContentLength() throws -> Int {
        var header = Data()
        let fh = stdoutPipe.fileHandleForReading

        while true {
            let byte = fh.readData(ofLength: 1)
            guard !byte.isEmpty else { throw LSPError.connectionClosed }
            header.append(byte)

            if header.count >= 4 {
                let tail = header.suffix(4)
                if tail == Data([0x0D, 0x0A, 0x0D, 0x0A]) { break }  // \r\n\r\n
            }
        }

        let headerStr = String(data: header, encoding: .utf8) ?? ""
        for line in headerStr.components(separatedBy: "\r\n") {
            if line.lowercased().hasPrefix("content-length:") {
                let value = line.dropFirst("content-length:".count).trimmingCharacters(in: .whitespaces)
                if let length = Int(value) { return length }
            }
        }
        throw LSPError.invalidHeader
    }
}