PadXcode-iPad/Network/LSPClient.swift

// LSPClient.swift
// CompletionItem, DefinitionLocation, DiagnosticRange all live in
// Shared/SharedModels.swift — not redeclared here.
import Foundation

// MARK: - LSP Request/Response types

struct LSPRequest: Encodable {
    let jsonrpc = "2.0"
    let id:     Int?
    let method: String
    let params: JSONValue?
}

struct LSPResponse: Decodable {
    let jsonrpc: String
    let id:      Int?
    let result:  JSONValue?
    let error:   LSPError?
    let method:  String?
    let params:  JSONValue?
}

struct LSPError: Decodable {
    let code:    Int
    let message: String
}

// MARK: - JSONValue (generic JSON tree for LSP params)

indirect enum JSONValue: Codable {
    case string(String)
    case int(Int)
    case double(Double)
    case bool(Bool)
    case object([String: JSONValue])
    case array([JSONValue])
    case null

    init(from decoder: Decoder) throws {
        let c = try decoder.singleValueContainer()
        if let v = try? c.decode(String.self)              { self = .string(v); return }
        if let v = try? c.decode(Int.self)                 { self = .int(v);    return }
        if let v = try? c.decode(Double.self)              { self = .double(v); return }
        if let v = try? c.decode(Bool.self)                { self = .bool(v);   return }
        if let v = try? c.decode([String: JSONValue].self) { self = .object(v); return }
        if let v = try? c.decode([JSONValue].self)         { self = .array(v);  return }
        self = .null
    }

    func encode(to encoder: Encoder) throws {
        var c = encoder.singleValueContainer()
        switch self {
        case .string(let v): try c.encode(v)
        case .int(let v):    try c.encode(v)
        case .double(let v): try c.encode(v)
        case .bool(let v):   try c.encode(v)
        case .object(let v): try c.encode(v)
        case .array(let v):  try c.encode(v)
        case .null:          try c.encodeNil()
        }
    }
}

// MARK: - LSP Client

@MainActor
final class LSPClient: ObservableObject {

    @Published var isConnected:         Bool = false
    @Published var completions:         [CompletionItem] = []
    @Published var diagnostics:         [String: [DiagnosticRange]] = [:]
    @Published var hoverText:           String?
    @Published var definitionLocation:  DefinitionLocation?

    private var webSocketTask:    URLSessionWebSocketTask?
    private var messageId         = 0
    private var pendingRequests:  [Int: CheckedContinuation<LSPResponse, Error>] = [:]
    var config:                   DaemonConfiguration

    init(config: DaemonConfiguration) { self.config = config }

    // MARK: - Connection

    func connect(projectPath: String) async {
        guard let startURL = URL(string: "\(config.baseURL)/lsp/start") else { return }
        var req = URLRequest(url: startURL)
        req.httpMethod = "POST"
        req.setValue("application/json", forHTTPHeaderField: "Content-Type")
        req.httpBody   = try? JSONEncoder().encode(["projectPath": projectPath])
        _ = try? await URLSession.shared.data(for: req)

        guard let wsURL = URL(string: config.lspWSURL) else { return }
        webSocketTask?.cancel()
        webSocketTask = URLSession.shared.webSocketTask(with: wsURL)
        webSocketTask?.resume()
        isConnected = true
        config.lspConnected = true

        await sendInitialize(projectPath: projectPath)
        Task { await receiveLoop() }
    }

    func disconnect() {
        webSocketTask?.cancel()
        webSocketTask   = nil
        isConnected     = false
        config.lspConnected = false
    }

    // MARK: - LSP Lifecycle

    private func sendInitialize(projectPath: String) async {
        let params = JSONValue.object([
            "processId": .int(Int(ProcessInfo.processInfo.processIdentifier)),
            "rootUri":   .string("file://\(projectPath)"),
            "capabilities": .object([
                "textDocument": .object([
                    "completion": .object(["completionItem": .object(["snippetSupport": .bool(false)])]),
                    "hover":      .object(["contentFormat": .array([.string("plaintext")])]),
                    "publishDiagnostics": .object([:])
                ])
            ]),
            "initializationOptions": .object([:])
        ])
        _ = try? await sendRequest(method: "initialize", params: params)
        sendNotification(method: "initialized", params: .object([:]))
    }

    // MARK: - Document Sync

    func didOpen(filePath: String, content: String) {
        sendNotification(method: "textDocument/didOpen", params: .object([
            "textDocument": .object([
                "uri":        .string("file://\(filePath)"),
                "languageId": .string(languageIdentifier(for: filePath)),
                "version":    .int(1),
                "text":       .string(content)
            ])
        ]))
    }

    func didChange(filePath: String, content: String, version: Int) {
        sendNotification(method: "textDocument/didChange", params: .object([
            "textDocument":   .object(["uri": .string("file://\(filePath)"), "version": .int(version)]),
            "contentChanges": .array([.object(["text": .string(content)])])
        ]))
    }

    func didSave(filePath: String) {
        sendNotification(method: "textDocument/didSave", params: .object([
            "textDocument": .object(["uri": .string("file://\(filePath)")])
        ]))
    }

    func didClose(filePath: String) {
        sendNotification(method: "textDocument/didClose", params: .object([
            "textDocument": .object(["uri": .string("file://\(filePath)")])
        ]))
    }

    // MARK: - Completions

    func requestCompletions(filePath: String, line: Int, column: Int) async -> [CompletionItem] {
        guard let response = try? await sendRequest(
            method: "textDocument/completion",
            params: textDocumentPositionParams(filePath: filePath, line: line - 1, character: column - 1)
        ) else { return [] }
        return parseCompletions(from: response)
    }

    // MARK: - Hover

    func requestHover(filePath: String, line: Int, column: Int) async -> String? {
        guard let response = try? await sendRequest(
            method: "textDocument/hover",
            params: textDocumentPositionParams(filePath: filePath, line: line - 1, character: column - 1)
        ) else { return nil }
        if case .object(let obj)      = response.result,
           case .object(let contents) = obj["contents"],
           case .string(let value)    = contents["value"] { return value }
        return nil
    }

    // MARK: - Definition

    func requestDefinition(filePath: String, line: Int, column: Int) async -> DefinitionLocation? {
        guard let response = try? await sendRequest(
            method: "textDocument/definition",
            params: textDocumentPositionParams(filePath: filePath, line: line - 1, character: column - 1)
        ) else { return nil }
        return parseDefinition(from: response)
    }

    // MARK: - WebSocket Send/Receive

    @discardableResult
    private func sendRequest(method: String, params: JSONValue?) async throws -> LSPResponse {
        messageId += 1
        let id = messageId
        let data = try JSONEncoder().encode(LSPRequest(id: id, method: method, params: params))
        guard let text = String(data: data, encoding: .utf8) else { throw URLError(.badServerResponse) }

        return try await withCheckedThrowingContinuation { cont in
            pendingRequests[id] = cont
            webSocketTask?.send(.string(text)) { error in
                if let error {
                    Task { @MainActor in
                        // Only resume if this request is still pending — handleMessage
                        // may have already received a response and resumed it.
                        if self.pendingRequests.removeValue(forKey: id) != nil {
                            cont.resume(throwing: error)
                        }
                    }
                }
            }
        }
    }

    private func sendNotification(method: String, params: JSONValue?) {
        guard let data = try? JSONEncoder().encode(LSPRequest(id: nil, method: method, params: params)),
              let text = String(data: data, encoding: .utf8) else { return }
        webSocketTask?.send(.string(text)) { _ in }
    }

    private func receiveLoop() async {
        guard let task = webSocketTask else {
            isConnected = false; return
        }
        while task.state == .running {
            do {
                let message = try await task.receive()
                switch message {
                case .string(let text): handleMessage(text)
                case .data(let data):
                    if let text = String(data: data, encoding: .utf8) { handleMessage(text) }
                @unknown default: break
                }
            } catch {
                isConnected = false
                config.lspConnected = false
                // Drain pending continuations so callers don't hang forever
                for (_, cont) in pendingRequests {
                    cont.resume(throwing: error)
                }
                pendingRequests.removeAll()
                return
            }
        }
        isConnected = false
    }

    private func handleMessage(_ text: String) {
        guard let data     = text.data(using: .utf8),
              let response = try? JSONDecoder().decode(LSPResponse.self, from: data) else { return }
        if let id = response.id, let cont = pendingRequests.removeValue(forKey: id) {
            cont.resume(returning: response); return
        }
        if let method = response.method { handleNotification(method: method, params: response.params) }
    }

    private func handleNotification(method: String, params: JSONValue?) {
        if method == "textDocument/publishDiagnostics" { handlePublishDiagnostics(params) }
    }

    private func handlePublishDiagnostics(_ params: JSONValue?) {
        guard case .object(let obj)  = params,
              case .string(let uri)  = obj["uri"],
              case .array(let diags) = obj["diagnostics"] else { return }

        let filePath = uri.replacingOccurrences(of: "file://", with: "")

        diagnostics[filePath] = diags.compactMap { diag -> DiagnosticRange? in
            guard case .object(let d)    = diag,
                  case .object(let rng)  = d["range"],
                  case .object(let start) = rng["start"],
                  case .int(let line)    = start["line"],
                  case .int(let col)     = start["character"],
                  case .int(let sev)     = d["severity"] else { return nil }
            let msg: String
            if case .string(let m) = d["message"] { msg = m } else { msg = "" }
            return DiagnosticRange(
                line: line + 1, column: col + 1, length: 30,
                severity: sev == 1 ? .error : .warning, message: msg
            )
        }
    }

    // MARK: - Parsing

    private func parseCompletions(from response: LSPResponse) -> [CompletionItem] {
        var items: [JSONValue] = []
        if case .array(let a) = response.result { items = a }
        else if case .object(let o) = response.result,
                case .array(let a)  = o["items"] { items = a }

        return items.compactMap { item -> CompletionItem? in
            guard case .object(let obj) = item,
                  case .string(let label) = obj["label"] else { return nil }
            let detail:     String; if case .string(let d) = obj["detail"]     { detail     = d } else { detail     = "" }
            let insertText: String?; if case .string(let t) = obj["insertText"] { insertText = t } else { insertText = nil }
            let filterText: String?; if case .string(let f) = obj["filterText"] { filterText = f } else { filterText = nil }
            let kind:       Int;     if case .int(let k)    = obj["kind"]        { kind       = k } else { kind       = 0 }
            return CompletionItem(label: label, detail: detail, kind: kind, insertText: insertText, filterText: filterText)
        }
    }

    private func parseDefinition(from response: LSPResponse) -> DefinitionLocation? {
        guard case .object(let loc) = response.result,
              case .string(let uri)  = loc["uri"],
              case .object(let rng)  = loc["range"],
              case .object(let start) = rng["start"],
              case .int(let line)    = start["line"],
              case .int(let col)     = start["character"] else { return nil }
        return DefinitionLocation(
            filePath: uri.replacingOccurrences(of: "file://", with: ""),
            line: line + 1, column: col + 1
        )
    }

    private func textDocumentPositionParams(filePath: String, line: Int, character: Int) -> JSONValue {
        .object([
            "textDocument": .object(["uri": .string("file://\(filePath)")]),
            "position":     .object(["line": .int(line), "character": .int(character)])
        ])
    }

    private func languageIdentifier(for path: String) -> String {
        switch (path as NSString).pathExtension.lowercased() {
        case "swift": return "swift"
        case "c":     return "c"
        case "cpp":   return "cpp"
        case "m":     return "objective-c"
        default:      return "plaintext"
        }
    }
}