//
// OpenRocky — Voice-first AI Agent
// https://github.com/openrocky
//
// Developed by everettjf with the assistance of Claude Code and Codex.
// Date: 2026-05-04
// Copyright (c) 2026 everettjf. All rights reserved.
//

import Foundation

/// Minimal MCP client implementing the Streamable-HTTP transport (the modern
/// transport — stdio is moot on iOS, and the legacy SSE GET path was deprecated
/// in 2025-06-18). Each `OpenRockyMCPClient` is short-lived: it spans the
/// initialize handshake plus one tools/list or tools/call. We re-init per call
/// to stay simple — MCP servers are typically stateless across the
/// initialize→call boundary, and any session token returned by the server is
/// carried through within the one client.
///
/// What this client deliberately does NOT do:
/// - resources, prompts, sampling — out of scope; we only need callable tools
/// - long-lived SSE streaming back-channel — only the inline response stream
/// - OAuth dynamic client registration — bring-your-own bearer token only
actor OpenRockyMCPClient {
    /// Mirrors the MCP protocol revision we target. Servers that advertise an
    /// older version on initialize will still typically work for tools/list
    /// and tools/call — the spec keeps these stable.
    static let protocolVersion = "2025-06-18"

    private let server: OpenRockyMCPServer
    private let urlSession: URLSession
    private let decoder = JSONDecoder()
    private let encoder = JSONEncoder()
    private var nextRequestID: Int = 1
    private var sessionID: String?

    init(server: OpenRockyMCPServer, urlSession: URLSession = .shared) {
        self.server = server
        self.urlSession = urlSession
    }

    // MARK: - Public API

    /// Run the initialize → tools/list dance and return the advertised tool
    /// catalog as the cache-shape struct the rest of the runtime understands.
    func listTools() async throws -> [OpenRockyMCPServer.CachedTool] {
        try await initializeIfNeeded()

        let response: ToolsListResult = try await send(method: "tools/list", params: nil)
        return response.tools.map { remote in
            OpenRockyMCPServer.CachedTool(
                name: remote.name,
                description: remote.description ?? "",
                inputSchemaJSON: Self.encodeAnyJSON(remote.inputSchema) ?? "{}"
            )
        }
    }

    /// Invoke a tool by name, with model-supplied JSON arguments. Returns the
    /// concatenated text content from the MCP response, or a JSON-encoded
    /// dump of the structured content if no text was provided.
    func callTool(name: String, argumentsJSON: String) async throws -> String {
        try await initializeIfNeeded()

        let argumentsValue = Self.parseAnyJSON(argumentsJSON) ?? AnyJSON.object([:])
        let params = AnyJSON.object([
            "name": .string(name),
            "arguments": argumentsValue,
        ])

        let result: ToolsCallResult = try await send(method: "tools/call", params: params)
        if result.isError == true {
            // Surface the model-visible failure as a thrown error so the
            // subagent loop records `succeeded: false` and the user sees the
            // ✗ marker in the timeline.
            throw OpenRockyMCPClientError.toolError(message: result.flattenedText())
        }
        let text = result.flattenedText()
        return text.isEmpty ? "{}" : text
    }

    // MARK: - Initialize

    private func initializeIfNeeded() async throws {
        guard sessionID == nil else { return }
        let params = AnyJSON.object([
            "protocolVersion": .string(Self.protocolVersion),
            "capabilities": .object([:]),
            "clientInfo": .object([
                "name": .string("OpenRocky"),
                "version": .string("1.0"),
            ]),
        ])
        // Many servers reject initialize being decoded into a typed struct
        // because the result contains free-form server capabilities. We
        // throw it through `AnyJSON` and don't inspect — we only care that
        // it succeeded and (optionally) returned a session id header.
        let _: AnyJSON = try await send(method: "initialize", params: params)
    }

    // MARK: - Transport

    /// Encode + POST a JSON-RPC request, decode the matching response. Handles
    /// both `application/json` and `text/event-stream` response bodies (the
    /// two encodings the Streamable-HTTP transport supports).
    private func send<R: Decodable & Sendable>(method: String, params: AnyJSON?) async throws -> R {
        let id = nextRequestID
        nextRequestID += 1

        let request = JSONRPCRequest(jsonrpc: "2.0", id: id, method: method, params: params)
        guard let url = URL(string: server.endpointURL) else {
            throw OpenRockyMCPClientError.invalidURL(server.endpointURL)
        }

        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = "POST"
        urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
        urlRequest.setValue("application/json, text/event-stream", forHTTPHeaderField: "Accept")
        urlRequest.setValue("OpenRocky/1.0", forHTTPHeaderField: "User-Agent")
        if let token = server.bearerToken, !token.isEmpty {
            urlRequest.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        }
        for (key, value) in server.extraHeaders {
            urlRequest.setValue(value, forHTTPHeaderField: key)
        }
        if let session = sessionID {
            urlRequest.setValue(session, forHTTPHeaderField: "Mcp-Session-Id")
        }
        urlRequest.httpBody = try encoder.encode(request)
        urlRequest.timeoutInterval = 30

        let (data, response) = try await urlSession.data(for: urlRequest)
        guard let http = response as? HTTPURLResponse else {
            throw OpenRockyMCPClientError.unexpectedResponse("not http")
        }

        // Capture session id if the server set one on this round-trip — we
        // need to echo it on subsequent calls within the same client.
        if let newSession = http.value(forHTTPHeaderField: "Mcp-Session-Id"), !newSession.isEmpty {
            sessionID = newSession
        }

        guard (200...299).contains(http.statusCode) else {
            let body = String(data: data, encoding: .utf8) ?? "<non-utf8>"
            throw OpenRockyMCPClientError.httpError(status: http.statusCode, body: String(body.prefix(400)))
        }

        // 202 Accepted with no body is the spec'd response for plain
        // notifications — no result expected. Our caller's R should be
        // AnyJSON in that case (we only use 202 for notifications, but
        // belt-and-suspenders: bail with a useful error).
        if http.statusCode == 202 {
            if let any = AnyJSON.null as? R { return any }
            throw OpenRockyMCPClientError.unexpectedResponse("202 with no body")
        }

        let contentType = (http.value(forHTTPHeaderField: "Content-Type") ?? "").lowercased()
        let envelope: JSONRPCResponse<R>
        if contentType.contains("text/event-stream") {
            envelope = try Self.decodeFirstSSE(data: data, decoder: decoder, expectingID: id)
        } else {
            envelope = try decoder.decode(JSONRPCResponse<R>.self, from: data)
        }

        if let error = envelope.error {
            throw OpenRockyMCPClientError.rpcError(code: error.code, message: error.message)
        }
        guard let result = envelope.result else {
            throw OpenRockyMCPClientError.unexpectedResponse("missing result")
        }
        return result
    }

    /// Streamable-HTTP can return SSE: a stream of `event: ...` / `data: ...`
    /// blocks, separated by blank lines. The first `data:` payload that
    /// matches our request id is the JSON-RPC response — we don't need any
    /// later events for non-streaming methods.
    private static func decodeFirstSSE<R: Decodable & Sendable>(
        data: Data,
        decoder: JSONDecoder,
        expectingID id: Int
    ) throws -> JSONRPCResponse<R> {
        guard let body = String(data: data, encoding: .utf8) else {
            throw OpenRockyMCPClientError.unexpectedResponse("non-utf8 sse body")
        }
        let lines = body.split(whereSeparator: { $0.isNewline })
        for line in lines {
            let trimmed = line.trimmingCharacters(in: .whitespaces)
            guard trimmed.hasPrefix("data:") else { continue }
            let payload = String(trimmed.dropFirst("data:".count)).trimmingCharacters(in: .whitespaces)
            guard !payload.isEmpty, payload != "[DONE]" else { continue }
            guard let payloadData = payload.data(using: .utf8) else { continue }
            // Decode as the typed envelope; if id matches, return — otherwise
            // continue scanning (servers can intersperse pings or progress).
            if let decoded = try? decoder.decode(JSONRPCResponse<R>.self, from: payloadData),
               decoded.id == id || decoded.error != nil {
                return decoded
            }
        }
        throw OpenRockyMCPClientError.unexpectedResponse("no matching JSON-RPC data event in SSE stream")
    }

    // MARK: - JSON helpers

    private static func encodeAnyJSON(_ value: AnyJSON?) -> String? {
        guard let value else { return nil }
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.sortedKeys]
        guard let data = try? encoder.encode(value) else { return nil }
        return String(data: data, encoding: .utf8)
    }

    private static func parseAnyJSON(_ source: String) -> AnyJSON? {
        guard let data = source.data(using: .utf8) else { return nil }
        return try? JSONDecoder().decode(AnyJSON.self, from: data)
    }
}

// MARK: - JSON-RPC envelopes
//
// All transport-layer types are explicitly `nonisolated` because the project
// defaults to `SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor`, which would
// otherwise stamp the Codable conformances as main-actor-bound and break
// their use from inside `actor OpenRockyMCPClient`.

nonisolated private struct JSONRPCRequest: Encodable, Sendable {
    let jsonrpc: String
    let id: Int
    let method: String
    let params: AnyJSON?
}

nonisolated private struct JSONRPCResponse<R: Decodable & Sendable>: Decodable, Sendable {
    let jsonrpc: String?
    let id: Int?
    let result: R?
    let error: JSONRPCError?
}

nonisolated private struct JSONRPCError: Decodable, Sendable {
    let code: Int
    let message: String
}

// MARK: - tools/list, tools/call result types

nonisolated struct ToolsListResult: Decodable, Sendable {
    let tools: [Tool]

    nonisolated struct Tool: Decodable, Sendable {
        let name: String
        let description: String?
        let inputSchema: AnyJSON?
    }
}

nonisolated struct ToolsCallResult: Decodable, Sendable {
    let content: [ContentBlock]?
    let isError: Bool?

    nonisolated struct ContentBlock: Decodable, Sendable {
        let type: String
        let text: String?
    }

    /// MCP responses are an array of typed content blocks. For our use the
    /// chat model only consumes text, so we concatenate the text blocks.
    /// Non-text blocks (images, embedded resources) are dropped — the chat
    /// model wouldn't know what to do with them at this layer anyway.
    func flattenedText() -> String {
        guard let content else { return "" }
        let texts = content.compactMap { $0.type == "text" ? $0.text : nil }
        return texts.joined(separator: "\n")
    }
}

// MARK: - AnyJSON

/// Schema-less JSON value, used for tool inputSchemas (whose shape we don't
/// model locally) and tool-call arguments (which the chat model emits as
/// arbitrary JSON). Codable in both directions so we can pass it through
/// JSON-RPC without losing fidelity.
nonisolated enum AnyJSON: Codable, Equatable, Sendable {
    case null
    case bool(Bool)
    case int(Int)
    case double(Double)
    case string(String)
    case array([AnyJSON])
    case object([String: AnyJSON])

    init(from decoder: Decoder) throws {
        let container = try decoder.singleValueContainer()
        if container.decodeNil() {
            self = .null
        } else if let value = try? container.decode(Bool.self) {
            self = .bool(value)
        } else if let value = try? container.decode(Int.self) {
            self = .int(value)
        } else if let value = try? container.decode(Double.self) {
            self = .double(value)
        } else if let value = try? container.decode(String.self) {
            self = .string(value)
        } else if let value = try? container.decode([AnyJSON].self) {
            self = .array(value)
        } else if let value = try? container.decode([String: AnyJSON].self) {
            self = .object(value)
        } else {
            throw DecodingError.dataCorruptedError(in: container, debugDescription: "Unrepresentable JSON value")
        }
    }

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

// MARK: - Errors

enum OpenRockyMCPClientError: LocalizedError {
    case invalidURL(String)
    case httpError(status: Int, body: String)
    case rpcError(code: Int, message: String)
    case unexpectedResponse(String)
    case toolError(message: String)

    var errorDescription: String? {
        switch self {
        case .invalidURL(let raw):
            return "Invalid MCP endpoint URL: \(raw)"
        case .httpError(let status, let body):
            return "MCP HTTP \(status): \(body)"
        case .rpcError(_, let message):
            return "MCP server error: \(message)"
        case .unexpectedResponse(let detail):
            return "Unexpected MCP response: \(detail)"
        case .toolError(let message):
            return message.isEmpty ? "MCP tool reported an error." : message
        }
    }
}
