CxLLM-IOS/CxLLMStudio/Services/MCPService.swift

// CxLLM Studio — iOS Application
// Services/MCPService.swift — Model Context Protocol service

import Foundation
import CxCode

@Observable
final class MCPService {

    // MARK: - State

    private(set) var isInitialized = false
    private(set) var pingOk = false
    private(set) var tools: [McpToolDef] = []
    private(set) var resources: [McpResource] = []
    private(set) var error: String?
    private(set) var statusMessage: String = "Not initialized"
    private(set) var isExecutingTool = false
    private(set) var isReadingResource = false
    private(set) var history: [MCPHistoryEntry] = []
    private(set) var totalToolCalls: Int = 0
    private(set) var totalResourceReads: Int = 0
    private(set) var lastInitialized: Date?

    private let gateway: CxGateway

    init(gateway: CxGateway) {
        self.gateway = gateway
    }

    // MARK: - Lifecycle

    func initialize() async {
        do {
            _ = try await gateway.mcp.initialize()
            isInitialized = true
            error = nil
            lastInitialized = Date()

            if let _ = try? await gateway.mcp.ping() {
                pingOk = true
            }

            tools = try await gateway.mcp.listTools()
            resources = try await gateway.mcp.listResources()
            statusMessage = "Ready — \(tools.count) tools, \(resources.count) resources"
        } catch {
            isInitialized = false
            self.error = error.localizedDescription
            statusMessage = "Error: \(error.localizedDescription)"
        }
    }

    func ping() async -> Bool {
        do {
            let _ = try await gateway.mcp.ping()
            pingOk = true
            statusMessage = "Pong OK"
            error = nil
            return true
        } catch {
            pingOk = false
            statusMessage = "Ping failed"
            self.error = error.localizedDescription
            return false
        }
    }

    func reload() async {
        do {
            tools = try await gateway.mcp.listTools()
            resources = try await gateway.mcp.listResources()
            statusMessage = "Reloaded — \(tools.count) tools, \(resources.count) resources"
            error = nil
        } catch {
            self.error = error.localizedDescription
            statusMessage = "Reload failed: \(error.localizedDescription)"
        }
    }

    // MARK: - Tool Execution

    struct ToolCallResult {
        let content: String
        let isError: Bool
    }

    func callTool(name: String, arguments: [String: Any]) async -> ToolCallResult {
        isExecutingTool = true

        do {
            let result = try await gateway.mcp.callTool(name: name, arguments: arguments)
            let content = result.content?.compactMap(\.text).joined(separator: "\n") ?? "No content"
            let isErr = result.isError == true
            let displayContent = isErr ? "[ERROR] \(content)" : content

            totalToolCalls += 1
            history.insert(
                MCPHistoryEntry(type: .tool, name: name, result: displayContent, time: Date(), success: !isErr),
                at: 0
            )
            isExecutingTool = false
            error = nil
            return ToolCallResult(content: displayContent, isError: isErr)
        } catch {
            let errMsg = "Error: \(error.localizedDescription)"
            history.insert(
                MCPHistoryEntry(type: .tool, name: name, result: errMsg, time: Date(), success: false),
                at: 0
            )
            self.error = error.localizedDescription
            isExecutingTool = false
            return ToolCallResult(content: errMsg, isError: true)
        }
    }

    // MARK: - Resource Reading

    func readResource(uri: String, name: String) async -> String {
        isReadingResource = true

        do {
            let content = try await gateway.mcp.readResource(uri: uri)
            totalResourceReads += 1
            history.insert(
                MCPHistoryEntry(type: .resource, name: name, result: String(content.prefix(100)), time: Date(), success: true),
                at: 0
            )
            isReadingResource = false
            error = nil
            return content
        } catch {
            let errMsg = "Error: \(error.localizedDescription)"
            history.insert(
                MCPHistoryEntry(type: .resource, name: name, result: errMsg, time: Date(), success: false),
                at: 0
            )
            self.error = error.localizedDescription
            isReadingResource = false
            return errMsg
        }
    }

    // MARK: - History Management

    func clearHistory() {
        history.removeAll()
    }

    // MARK: - Computed

    var connectionState: ConnectionState {
        if !isInitialized { return .disconnected }
        if pingOk { return .live }
        return .initialized
    }

    var successRate: Double {
        guard !history.isEmpty else { return 0 }
        let successes = history.filter(\.success).count
        return Double(successes) / Double(history.count)
    }

    func tool(named name: String) -> McpToolDef? {
        tools.first { $0.name == name }
    }

    func filteredTools(query: String) -> [McpToolDef] {
        guard !query.isEmpty else { return tools }
        let q = query.lowercased()
        return tools.filter {
            $0.name.lowercased().contains(q) || ($0.description ?? "").lowercased().contains(q)
        }
    }

    func filteredResources(query: String) -> [McpResource] {
        guard !query.isEmpty else { return resources }
        let q = query.lowercased()
        return resources.filter {
            $0.name.lowercased().contains(q) || $0.uri.lowercased().contains(q)
        }
    }
}

// MARK: - Supporting Types

enum ConnectionState {
    case disconnected, initialized, live

    var label: String {
        switch self {
        case .disconnected: return "Off"
        case .initialized:  return "Init"
        case .live:         return "Live"
        }
    }

    var isConnected: Bool {
        self != .disconnected
    }
}

struct MCPHistoryEntry: Identifiable {
    let id = UUID()
    let type: EntryType
    let name: String
    let result: String
    let time: Date
    let success: Bool

    enum EntryType: String {
        case tool, resource
    }
}