Termind/TFI/Application/FeatureModels/MobileMCPEditorModel.swift

import AgentKit
import Foundation
import Observation
import TermindApplication

@MainActor @Observable
final class MobileMCPEditorModel {
    typealias Discover = @Sendable (MCPServer, String?) async throws -> MCPDiscovery
    let original: MCPServer
    var draft: MCPServer
    let key = MobileSecretDraft()
    private(set) var running = false
    private(set) var connection: String?
    private let ai: AIConfigurationModel
    private let wasPersisted: Bool
    private let discover: Discover
    private var task: Task<Void, Never>?
    private var generation = UUID()
    private var policies: [String: AgentToolDescriptor.ApprovalPolicy] = [:]

    init(
        ai: AIConfigurationModel, server: MCPServer,
        discover: @escaping Discover = {
            try await MCPClient(
                clientInfo: .init(
                    name: "Termind",
                    version: Bundle.main.infoDictionary?["CFBundleShortVersionString"] as? String ?? "0")
            ).discover($0, bearerToken: $1)
        }
    ) {
        self.ai = ai
        original = server
        draft = server
        wasPersisted = ai.mcpServer(server.id) != nil
        self.discover = discover
        key.begin(in: ai, reference: server.credentialRef)
    }
    var hasValidURL: Bool {
        guard let url = URL(string: draft.url.trimmingCharacters(in: .whitespacesAndNewlines)),
            let scheme = url.scheme?.lowercased(), let host = url.host, !host.isEmpty
        else { return false }
        return scheme == "https" || scheme == "http"
    }
    var canSave: Bool {
        !draft.name.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty && !draft.effectiveNamespaceID.isEmpty
            && !ai.hasMCPNamespaceConflict(draft.effectiveNamespaceID, excluding: draft.id)
            && hasValidURL && key.isValid
    }
    func connect() {
        cancel()
        guard canSave, !wasPersisted || ai.mcpServer(original.id) != nil else { return }
        let secret: String
        do {
            secret = try key.effective(
                in: ai, reference: ai.mcpServer(original.id)?.credentialRef ?? draft.credentialRef)
        } catch {
            connection = error.localizedDescription
            return
        }
        for tool in original.tools + draft.tools { policies[tool.id] = tool.accessPolicy }
        let server = draft
        let persisted = ai.mcpServer(original.id)
        let keyAction = key.action
        let keyText = key.text
        let token = generation
        running = true
        task = Task { [weak self, discover] in
            defer { if self?.generation == token { self?.running = false } }
            do {
                let outcome = try await discover(server, secret.isEmpty ? nil : secret)
                guard let self, generation == token, !Task.isCancelled,
                    ai.mcpServer(original.id) == persisted, draft.name == server.name,
                    draft.url == server.url, draft.transport == server.transport, draft.headers == server.headers,
                    key.action == keyAction, key.text == keyText
                else { return }
                for tool in draft.tools { policies[tool.id] = tool.accessPolicy }
                var seen = Set<String>()
                draft.tools = outcome.tools.filter { seen.insert($0.id).inserted }.map {
                    var tool = $0
                    tool.accessPolicy = policies[tool.id] ?? .ask
                    tool.enforceTermindApprovalPolicy()
                    return tool
                }
                draft.negotiatedVersion = outcome.protocolVersion
                draft.lastConnectedAt = .now
                connection = String(localized: "\(outcome.serverName) — \(draft.tools.count) tools")
            } catch {
                guard let self, generation == token, !Task.isCancelled,
                    ai.mcpServer(original.id) == persisted, draft.name == server.name,
                    draft.url == server.url, draft.transport == server.transport, draft.headers == server.headers,
                    key.action == keyAction, key.text == keyText
                else { return }
                connection = error.localizedDescription
            }
        }
    }
    func cancel() {
        generation = UUID()
        task?.cancel()
        task = nil
        running = false
        connection = nil
    }
    func invalidateConnection() {
        cancel()
        for tool in draft.tools { policies[tool.id] = tool.accessPolicy }
        draft.tools = []
        draft.negotiatedVersion = ""
        draft.lastConnectedAt = nil
    }

    func upsertHeader(_ header: MCPHeader) {
        if let index = draft.headers.firstIndex(where: { $0.id == header.id }) {
            guard draft.headers[index] != header else { return }
            draft.headers[index] = header
        } else {
            draft.headers.append(header)
        }
        invalidateConnection()
    }

    func removeHeaders(atOffsets offsets: IndexSet) {
        guard !offsets.isEmpty else { return }
        for index in offsets.sorted(by: >) { draft.headers.remove(at: index) }
        invalidateConnection()
    }

    func save() throws {
        guard canSave else { return }
        var value = draft
        if wasPersisted {
            guard let current = ai.mcpServer(original.id) else {
                throw MobileEditorError.message(String(localized: "Server Deleted"))
            }
            var fields = SettingsFields(original, draft, current)
            try fields.merge(\.name)
            try fields.merge(\.transport)
            try fields.merge(\.url)
            try fields.merge(\.isEnabled)
            value = fields.result
            value.headers = try SettingsDraftMerge.collection(original.headers, draft.headers, current.headers) {
                old, local, remote in
                var fields = SettingsFields(old, local, remote)
                try fields.merge(\.name)
                try fields.merge(\.value)
                return fields.result
            }
            value.tools = try SettingsDraftMerge.collection(original.tools, draft.tools, current.tools) {
                old, local, remote in
                var fields = SettingsFields(old, local, remote)
                try fields.merge(\.accessPolicy)
                try fields.merge(\.title)
                try fields.merge(\.summary)
                try fields.merge(\.inputSchema)
                try fields.merge(\.annotations)
                return fields.result
            }
            if draft.lastConnectedAt != original.lastConnectedAt {
                value.lastConnectedAt = draft.lastConnectedAt
                value.negotiatedVersion = draft.negotiatedVersion
            }
        }
        value.url = value.url.trimmingCharacters(in: .whitespacesAndNewlines)
        if value.namespaceID.isEmpty { value.namespaceID = value.effectiveNamespaceID }
        guard !ai.hasMCPNamespaceConflict(value.effectiveNamespaceID, excluding: value.id) else {
            throw MobileEditorError.message(String(localized: "Another MCP server already uses this ID."))
        }
        try key.commit(in: ai, reference: value.credentialRef, kind: "mcp") { reference in
            value.credentialRef = reference
            ai.upsert(value)
            return ai.storageError == nil
        }
    }
}