oai-swift/oAI/Services/AnytypeMCPService.swift

//
//  AnytypeMCPService.swift
//  oAI
//
//  Anytype MCP integration via local HTTP API at http://127.0.0.1:31009
//
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen
//
// This file is part of oAI.
//
// oAI is licensed under the PolyForm Noncommercial License 1.0.0.
// You may use, study, modify, and share it for any noncommercial
// purpose. Commercial use — including selling oAI or any part of
// it, standalone or bundled into another product or service —
// requires a separate commercial license from the copyright holder.
//
// See the LICENSE file or
// <https://polyformproject.org/licenses/noncommercial/1.0.0> for
// the full license text. For commercial licensing, contact Rune
// Olsen via <https://oai.pm>.


import Foundation
import os

@Observable
class AnytypeMCPService {
    static let shared = AnytypeMCPService()

    private let settings = SettingsService.shared
    private let log = Logger(subsystem: "com.oai.oAI", category: "mcp")
    private let apiVersion = "2025-11-08"
    private let timeout: TimeInterval = 10

    private(set) var isConnected = false

    private init() {}

    // MARK: - Connection Test

    func testConnection() async -> Result<String, Error> {
        do {
            let result = try await request(endpoint: "/v1/spaces", method: "GET", body: nil)
            if let spaces = result["data"] as? [[String: Any]] {
                isConnected = true
                return .success("Connected (\(spaces.count) space\(spaces.count == 1 ? "" : "s"))")
            } else {
                isConnected = true
                return .success("Connected to Anytype")
            }
        } catch {
            isConnected = false
            return .failure(error)
        }
    }

    // MARK: - Tool Schemas

    func getToolSchemas() -> [Tool] {
        return [
            makeTool(
                name: "anytype_search_global",
                description: "Search across all Anytype spaces for objects matching a query. Returns matching objects with their IDs, names, types, and space info.",
                properties: [
                    "query": prop("string", "The search query text"),
                    "limit": prop("number", "Maximum number of results to return (default: 20)")
                ],
                required: ["query"]
            ),
            makeTool(
                name: "anytype_list_spaces",
                description: "List all available Anytype spaces (workspaces). Returns space IDs, names, and basic info.",
                properties: [:],
                required: []
            ),
            makeTool(
                name: "anytype_get_space_objects",
                description: "Get objects in a specific Anytype space. Returns a list of objects with their IDs, names, and types. Use search tools to find specific objects rather than listing all.",
                properties: [
                    "space_id": prop("string", "The ID of the Anytype space"),
                    "limit": prop("number", "Maximum number of objects to return (default: 20, max recommended: 50)")
                ],
                required: ["space_id"]
            ),
            makeTool(
                name: "anytype_get_object",
                description: "Get the full details and content of a specific Anytype object by ID. The 'body' field contains the COMPLETE markdown including Anytype internal links (anytype://object?...) and file links that MUST be preserved exactly when updating.",
                properties: [
                    "space_id": prop("string", "The ID of the space containing the object"),
                    "object_id": prop("string", "The ID of the object to retrieve")
                ],
                required: ["space_id", "object_id"]
            ),
            makeTool(
                name: "anytype_create_object",
                description: "Create a new object (note, task, or page) in an Anytype space.",
                properties: [
                    "space_id": prop("string", "The ID of the space to create the object in"),
                    "name": prop("string", "The name/title of the object"),
                    "body": prop("string", "The text content/body of the object"),
                    "type": prop("string", "The type of object: 'note', 'task', or 'page' (default: note)")
                ],
                required: ["space_id", "name"]
            ),
            makeTool(
                name: "anytype_append_to_object",
                description: """
                Append new markdown content to the end of an existing Anytype object without touching the existing body. \
                This is the PREFERRED way to add content to existing notes, pages, or tasks — \
                it preserves all Anytype internal links (anytype://...) and mention blocks exactly. \
                Use this instead of anytype_update_object whenever you are adding information rather than rewriting.
                """,
                properties: [
                    "space_id":  prop("string", "The ID of the space containing the object"),
                    "object_id": prop("string", "The ID of the object to append to"),
                    "content":   prop("string", "Markdown content to append at the end of the object body")
                ],
                required: ["space_id", "object_id", "content"]
            ),
            makeTool(
                name: "anytype_update_object",
                description: """
                Replace the full markdown body or rename an Anytype object. \
                WARNING: This replaces the ENTIRE body — prefer anytype_append_to_object for adding content \
                to existing objects, as full replacement degrades rich Anytype internal links (anytype://...) to plain text. \
                Use this ONLY when you truly need to rewrite or restructure existing content. \
                CRITICAL RULES when using this tool: \
                1) Always call anytype_get_object first to get the current EXACT markdown. \
                2) Make ONLY the minimal requested change — nothing else. \
                3) Copy ALL other content CHARACTER-FOR-CHARACTER, including anytype://object?objectId=... links, \
                file links, headings, and all formatting. Do NOT summarise, paraphrase, or reformat anything you are not asked to change.
                """,
                properties: [
                    "space_id": prop("string", "The ID of the space containing the object"),
                    "object_id": prop("string", "The ID of the object to update"),
                    "name": prop("string", "New name/title for the object (optional)"),
                    "body": prop("string", "COMPLETE replacement markdown content — must preserve ALL existing content with ONLY the minimal intended change applied")
                ],
                required: ["space_id", "object_id"]
            ),
            makeTool(
                name: "anytype_set_done",
                description: "Mark a task as done or undone in Anytype by setting its 'done' relation. This is the correct way to check/uncheck tasks — do not use anytype_update_object body text for this.",
                properties: [
                    "space_id": prop("string", "The ID of the space containing the task"),
                    "object_id": prop("string", "The ID of the task to update"),
                    "done": prop("boolean", "true to mark as done, false to mark as undone")
                ],
                required: ["space_id", "object_id", "done"]
            ),
            makeTool(
                name: "anytype_toggle_checkbox",
                description: """
                Surgically toggle a specific checkbox in an Anytype object without rewriting the full body. \
                Use this INSTEAD of anytype_update_object when marking a to-do item done/undone. \
                Provide partial text that uniquely identifies the checkbox line (do not include the '- [ ]' or '- [x]' prefix). \
                The tool finds the matching line and flips its checkbox state, preserving all other content exactly.
                """,
                properties: [
                    "space_id": prop("string", "The ID of the space containing the object"),
                    "object_id": prop("string", "The ID of the object containing the checkbox"),
                    "line_text": prop("string", "Partial text of the checkbox line to find (e.g. 'Buy groceries'). Do NOT include the '- [ ]' prefix."),
                    "done": prop("boolean", "true to check the box, false to uncheck it")
                ],
                required: ["space_id", "object_id", "line_text", "done"]
            ),
            makeTool(
                name: "anytype_search_space",
                description: "Search within a specific Anytype space for objects matching a query.",
                properties: [
                    "space_id": prop("string", "The ID of the Anytype space to search in"),
                    "query": prop("string", "The search query text"),
                    "limit": prop("number", "Maximum number of results (default: 20)")
                ],
                required: ["space_id", "query"]
            )
        ]
    }

    // MARK: - Tool Execution

    func executeTool(name: String, arguments: String) async -> [String: Any] {
        log.info("Executing Anytype tool: \(name)")

        guard let argData = arguments.data(using: .utf8),
              let args = try? JSONSerialization.jsonObject(with: argData) as? [String: Any] else {
            return ["error": "Invalid arguments JSON"]
        }

        return await executeToolAsync(name: name, args: args)
    }

    private func executeToolAsync(name: String, args: [String: Any]) async -> [String: Any] {
        do {
            switch name {
            case "anytype_search_global":
                guard let query = args["query"] as? String else {
                    return ["error": "Missing required parameter: query"]
                }
                let limit = args["limit"] as? Int ?? 20
                return try await searchGlobal(query: query, limit: limit)

            case "anytype_list_spaces":
                return try await listSpaces()

            case "anytype_get_space_objects":
                guard let spaceId = args["space_id"] as? String else {
                    return ["error": "Missing required parameter: space_id"]
                }
                let limit = min(args["limit"] as? Int ?? 20, 50)
                return try await getSpaceObjects(spaceId: spaceId, limit: limit)

            case "anytype_get_object":
                guard let spaceId = args["space_id"] as? String,
                      let objectId = args["object_id"] as? String else {
                    return ["error": "Missing required parameters: space_id, object_id"]
                }
                return try await getObject(spaceId: spaceId, objectId: objectId)

            case "anytype_create_object":
                guard let spaceId = args["space_id"] as? String,
                      let name = args["name"] as? String else {
                    return ["error": "Missing required parameters: space_id, name"]
                }
                let body = args["body"] as? String ?? ""
                let type_ = args["type"] as? String ?? "note"
                return try await createObject(spaceId: spaceId, name: name, body: body, type: type_)

            case "anytype_append_to_object":
                guard let spaceId = args["space_id"] as? String,
                      let objectId = args["object_id"] as? String,
                      let content = args["content"] as? String else {
                    return ["error": "Missing required parameters: space_id, object_id, content"]
                }
                return try await appendToObject(spaceId: spaceId, objectId: objectId, content: content)

            case "anytype_update_object":
                guard let spaceId = args["space_id"] as? String,
                      let objectId = args["object_id"] as? String else {
                    return ["error": "Missing required parameters: space_id, object_id"]
                }
                let name = args["name"] as? String
                let body = args["body"] as? String
                return try await updateObject(spaceId: spaceId, objectId: objectId, name: name, body: body)

            case "anytype_set_done":
                guard let spaceId = args["space_id"] as? String,
                      let objectId = args["object_id"] as? String else {
                    return ["error": "Missing required parameters: space_id, object_id"]
                }
                // Accept done as Bool or as Int (1/0) since JSON can arrive either way
                let done: Bool
                if let b = args["done"] as? Bool {
                    done = b
                } else if let n = args["done"] as? Int {
                    done = n != 0
                } else {
                    return ["error": "Missing or invalid parameter: done (expected boolean)"]
                }
                return try await setDone(spaceId: spaceId, objectId: objectId, done: done)

            case "anytype_toggle_checkbox":
                guard let spaceId = args["space_id"] as? String,
                      let objectId = args["object_id"] as? String,
                      let lineText = args["line_text"] as? String else {
                    return ["error": "Missing required parameters: space_id, object_id, line_text"]
                }
                let done: Bool
                if let b = args["done"] as? Bool { done = b }
                else if let n = args["done"] as? Int { done = n != 0 }
                else { return ["error": "Missing or invalid parameter: done (expected boolean)"] }
                return try await toggleCheckbox(spaceId: spaceId, objectId: objectId, lineText: lineText, done: done)

            case "anytype_search_space":
                guard let spaceId = args["space_id"] as? String,
                      let query = args["query"] as? String else {
                    return ["error": "Missing required parameters: space_id, query"]
                }
                let limit = args["limit"] as? Int ?? 20
                return try await searchSpace(spaceId: spaceId, query: query, limit: limit)

            default:
                return ["error": "Unknown Anytype tool: \(name)"]
            }
        } catch AnytypeError.notRunning {
            return ["error": "Cannot connect to Anytype. Make sure the desktop app is running."]
        } catch AnytypeError.unauthorized {
            return ["error": "Invalid API key. Check your Anytype API key in Settings > MCP."]
        } catch AnytypeError.httpError(let code, let msg) {
            return ["error": "Anytype API error \(code): \(msg)"]
        } catch {
            return ["error": "Anytype error: \(error.localizedDescription)"]
        }
    }

    // MARK: - API Operations

    private func searchGlobal(query: String, limit: Int) async throws -> [String: Any] {
        let body: [String: Any] = ["query": query, "limit": limit]
        let result = try await request(endpoint: "/v1/search", method: "POST", body: body)

        if let objects = result["data"] as? [[String: Any]] {
            let formatted = objects.map { formatObject($0) }
            return ["count": formatted.count, "objects": formatted]
        }
        return ["count": 0, "objects": [], "message": "No results found"]
    }

    private func listSpaces() async throws -> [String: Any] {
        let result = try await request(endpoint: "/v1/spaces", method: "GET", body: nil)

        if let spaces = result["data"] as? [[String: Any]] {
            let formatted = spaces.map { space -> [String: Any] in
                [
                    "id": space["id"] ?? "",
                    "name": space["name"] ?? "Unnamed Space",
                    "type": space["spaceType"] ?? "unknown"
                ]
            }
            return ["count": formatted.count, "spaces": formatted]
        }
        return ["count": 0, "spaces": []]
    }

    private func getSpaceObjects(spaceId: String, limit: Int) async throws -> [String: Any] {
        let result = try await request(
            endpoint: "/v1/spaces/\(spaceId)/objects",
            method: "GET",
            body: nil,
            queryParams: ["limit": String(limit)]
        )

        if let objects = result["data"] as? [[String: Any]] {
            let formatted = Array(objects.prefix(limit).map { formatObject($0) })
            return ["count": formatted.count, "objects": formatted]
        }
        return ["count": 0, "objects": []]
    }

    private func getObject(spaceId: String, objectId: String) async throws -> [String: Any] {
        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)", method: "GET", body: nil)

        if let object = result["object"] as? [String: Any] {
            return formatObjectDetail(object)
        }
        // Return full raw response if expected structure not found
        return result
    }

    private func createObject(spaceId: String, name: String, body: String, type: String) async throws -> [String: Any] {
        let typeKey: String
        switch type.lowercased() {
        case "task": typeKey = "ot-task"
        case "page": typeKey = "ot-page"
        default: typeKey = "ot-note"
        }

        let requestBody: [String: Any] = [
            "name": name,
            "body": body,        // some API versions use "body"
            "markdown": body,    // newer API versions use "markdown"
            "typeKey": typeKey,
            "template_id": ""
        ]

        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/objects", method: "POST", body: requestBody)

        if let object = result["object"] as? [String: Any] {
            return ["success": true, "id": object["id"] ?? "", "name": name, "message": "Object created successfully"]
        }
        return ["success": true, "message": "Object created"]
    }

    private func appendToObject(spaceId: String, objectId: String, content: String) async throws -> [String: Any] {
        // Fetch current body
        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)", method: "GET", body: nil)
        guard let object = result["object"] as? [String: Any] else {
            return ["error": "Object not found"]
        }

        let existing: String
        if let md = object["markdown"] as? String { existing = md }
        else if let body = object["body"] as? String { existing = body }
        else { existing = "" }

        let separator = existing.isEmpty ? "" : "\n\n"
        let newMarkdown = existing + separator + content

        _ = try await request(
            endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)",
            method: "PATCH",
            body: ["markdown": newMarkdown]
        )
        return ["success": true, "message": "Content appended successfully"]
    }

    private func updateObject(spaceId: String, objectId: String, name: String?, body: String?) async throws -> [String: Any] {
        var requestBody: [String: Any] = [:]
        if let name = name { requestBody["name"] = name }
        // The Anytype API GET response uses "markdown" for body content, not "body"
        if let body = body { requestBody["markdown"] = body }

        guard !requestBody.isEmpty else {
            return ["error": "No fields to update. Provide name or body."]
        }

        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)", method: "PATCH", body: requestBody)
        if result.isEmpty {
            return ["success": true, "message": "Object updated (empty response from Anytype API)"]
        }
        return result
    }

    private func toggleCheckbox(spaceId: String, objectId: String, lineText: String, done: Bool) async throws -> [String: Any] {
        // Fetch current content
        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)", method: "GET", body: nil)
        guard let object = result["object"] as? [String: Any] else {
            return ["error": "Object not found"]
        }

        let markdown: String
        if let md = object["markdown"] as? String { markdown = md }
        else if let body = object["body"] as? String { markdown = body }
        else { return ["error": "Object has no markdown body"] }

        // Find the line containing the search text (case-insensitive)
        let lines = markdown.components(separatedBy: "\n")
        let searchLower = lineText.lowercased()
        guard let idx = lines.firstIndex(where: {
            let stripped = $0.replacingOccurrences(of: "- [ ] ", with: "")
                              .replacingOccurrences(of: "- [x] ", with: "")
                              .replacingOccurrences(of: "- [X] ", with: "")
            return stripped.lowercased().contains(searchLower)
        }) else {
            return ["error": "No checkbox line found containing: '\(lineText)'"]
        }

        var line = lines[idx]
        let wasDone = line.hasPrefix("- [x] ") || line.hasPrefix("- [X] ")

        if done == wasDone {
            return ["success": true, "message": "Checkbox was already \(done ? "checked" : "unchecked")", "line": line]
        }

        // Flip the checkbox marker, preserving the rest of the line exactly
        if done {
            if line.hasPrefix("- [ ] ") { line = "- [x] " + line.dropFirst(6) }
        } else {
            if line.hasPrefix("- [x] ") || line.hasPrefix("- [X] ") { line = "- [ ] " + line.dropFirst(6) }
        }

        var updated = lines
        updated[idx] = line
        let newMarkdown = updated.joined(separator: "\n")

        _ = try await request(
            endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)",
            method: "PATCH",
            body: ["markdown": newMarkdown]
        )
        return ["success": true, "done": done, "line": line, "message": "Checkbox updated"]
    }

    private func setDone(spaceId: String, objectId: String, done: Bool) async throws -> [String: Any] {
        // Anytype stores task completion as a relation called "done"
        let requestBody: [String: Any] = [
            "properties": [
                "done": done
            ]
        ]
        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/objects/\(objectId)", method: "PATCH", body: requestBody)
        if result.isEmpty {
            return ["success": true, "done": done, "message": "setDone sent (empty response from Anytype API — verify in app)"]
        }
        return result
    }

    private func searchSpace(spaceId: String, query: String, limit: Int) async throws -> [String: Any] {
        let body: [String: Any] = ["query": query, "limit": limit]
        let result = try await request(endpoint: "/v1/spaces/\(spaceId)/search", method: "POST", body: body)

        if let objects = result["data"] as? [[String: Any]] {
            let formatted = objects.map { formatObject($0) }
            return ["count": formatted.count, "objects": formatted]
        }
        return ["count": 0, "objects": [], "message": "No results found"]
    }

    // MARK: - HTTP Client

    private func request(endpoint: String, method: String, body: [String: Any]?, queryParams: [String: String] = [:]) async throws -> [String: Any] {
        guard let apiKey = settings.anytypeMcpAPIKey, !apiKey.isEmpty else {
            throw AnytypeError.unauthorized
        }

        let baseURL = settings.anytypeMcpEffectiveURL
        var urlString = baseURL + endpoint
        if !queryParams.isEmpty {
            var comps = URLComponents(string: urlString) ?? URLComponents()
            comps.queryItems = queryParams.map { URLQueryItem(name: $0.key, value: $0.value) }
            urlString = comps.url?.absoluteString ?? urlString
        }
        guard let url = URL(string: urlString) else {
            throw AnytypeError.httpError(0, "Invalid URL: \(urlString)")
        }

        var urlRequest = URLRequest(url: url, timeoutInterval: timeout)
        urlRequest.httpMethod = method
        urlRequest.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.setValue(apiVersion, forHTTPHeaderField: "Anytype-Version")

        if let body = body, method != "GET" {
            urlRequest.httpBody = try? JSONSerialization.data(withJSONObject: body)
        }

        do {
            let (data, response) = try await URLSession.shared.data(for: urlRequest)

            guard let httpResponse = response as? HTTPURLResponse else {
                throw AnytypeError.httpError(0, "Invalid response")
            }

            if httpResponse.statusCode == 401 {
                throw AnytypeError.unauthorized
            }

            guard (200...299).contains(httpResponse.statusCode) else {
                let msg = String(data: data, encoding: .utf8) ?? "Unknown error"
                throw AnytypeError.httpError(httpResponse.statusCode, msg)
            }

            guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
                return [:]
            }
            return json

        } catch let error as AnytypeError {
            throw error
        } catch {
            // Connection refused or network error = Anytype not running
            throw AnytypeError.notRunning
        }
    }

    // MARK: - Helpers

    private func formatObject(_ obj: [String: Any]) -> [String: Any] {
        [
            "id": obj["id"] ?? "",
            "name": obj["name"] ?? "Unnamed",
            "type": obj["type"] ?? obj["objectType"] ?? "unknown",
            "space_id": obj["spaceId"] ?? ""
        ]
    }

    private func formatObjectDetail(_ obj: [String: Any]) -> [String: Any] {
        var result: [String: Any] = formatObject(obj)
        // Anytype API returns content as "markdown" not "body"
        if let markdown = obj["markdown"] as? String { result["body"] = markdown }
        else if let body = obj["body"] as? String { result["body"] = body }
        if let snippet = obj["snippet"] as? String { result["snippet"] = snippet }
        return result
    }

    private func makeTool(name: String, description: String, properties: [String: Tool.Function.Parameters.Property], required: [String]) -> Tool {
        Tool(
            type: "function",
            function: Tool.Function(
                name: name,
                description: description,
                parameters: Tool.Function.Parameters(
                    type: "object",
                    properties: properties,
                    required: required
                )
            )
        )
    }

    private func prop(_ type: String, _ description: String) -> Tool.Function.Parameters.Property {
        Tool.Function.Parameters.Property(type: type, description: description, enum: nil)
    }
}

// MARK: - Error Types

enum AnytypeError: LocalizedError {
    case notRunning
    case unauthorized
    case httpError(Int, String)

    var errorDescription: String? {
        switch self {
        case .notRunning:
            return "Cannot connect to Anytype. Make sure the desktop app is running."
        case .unauthorized:
            return "Invalid API key. Check your Anytype API key in Settings > MCP."
        case .httpError(let code, let msg):
            return "Anytype API error \(code): \(msg)"
        }
    }
}