// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen

import Foundation

// MARK: - Transport abstraction

/// One MCP server connection's wire protocol. `ExternalMCPClient` builds JSON-RPC envelopes and
/// decodes typed results; the transport only owns *delivery* — how the bytes get to the server
/// and back. A long-lived subprocess pipe for stdio, discrete HTTP requests for Streamable HTTP.
@MainActor
protocol MCPTransport: AnyObject {
    /// Prepares the transport for use — launches the subprocess for stdio; a no-op for HTTP,
    /// since there's no persistent connection to establish ahead of the first request.
    func prepare() async throws
    /// Sends a JSON-RPC *request* (`message` includes `"id"`) and returns its `result` payload.
    /// Throws `MCPClientError.invalidResponse` if the server returned a JSON-RPC error, or
    /// `.timeout` if `timeoutSeconds` elapses first.
    func sendRequest(_ message: [String: Any], id: Int, timeoutSeconds: Double) async throws -> Data
    /// Sends a JSON-RPC *notification* (no `"id"`, no response expected).
    func sendNotification(_ message: [String: Any]) async throws
    /// Tears down the transport — terminates the subprocess / drops HTTP session state.
    func stop()
}

enum MCPTransportSupport {
    /// Extracts the `result` payload from a decoded top-level JSON-RPC response object, shared
    /// by every transport so error/result semantics stay identical regardless of wire protocol.
    static func extractResult(from json: [String: Any]) throws -> Data {
        if let err = json["error"] as? [String: Any] {
            throw MCPClientError.invalidResponse(err["message"] as? String ?? "Unknown error")
        }
        guard let result = json["result"],
              let resultData = try? JSONSerialization.data(withJSONObject: result)
        else {
            throw MCPClientError.invalidResponse("Missing result field")
        }
        return resultData
    }
}

// MARK: - Login Shell PATH Resolution

/// GUI apps launched from Finder/Dock/LaunchServices inherit launchd's minimal PATH
/// (`/usr/bin:/bin:/usr/sbin:/sbin`), not the interactive-shell PATH a Terminal session gets — so
/// `npx`/`node` installed via nvm, Homebrew, Volta etc. are invisible to a bare `/usr/bin/env <cmd>`
/// spawn, even though the same command works fine when the user runs it themselves in Terminal.
///
/// History: the obvious fix — spawn the user's real login shell and ask it for `$PATH` — caused two
/// real hangs in one session and was abandoned. First with `-ilc` (interactive login), which
/// deadlocked on a full pipe buffer from shell startup noise (fixed by draining concurrently).
/// After that fix, it hung *again* with non-interactive `-lc`: a live `sample` of the stuck process
/// showed `NSTask.waitUntilExit()` still parked inside a nested CFRunLoop even though the child
/// shell process was already gone from the process list — a CFRunLoop/SIGCHLD-notification
/// reentrancy issue specific to calling it from a Swift-concurrency worker thread inside this app's
/// XCTest host process. Spawning any subprocess just to read one environment variable is exposed to
/// this whole class of platform quirk, and a subprocess-based timeout can't protect against a hang
/// inside the *notification mechanism itself*. Instead: probe the fixed set of directories every
/// common Node install method actually uses, with plain filesystem checks — no subprocess, no run
/// loop, nothing that can hang, at the cost of being a known list rather than a fully general answer.
enum LoginShellEnvironment {
    /// Directories to add if present, in priority order. Covers Homebrew (Apple Silicon and Intel),
    /// MacPorts, Volta, and nvm's "default" alias (via `nvmDefaultNodeBinDirectory`).
    nonisolated static func candidateDirectories(home: String = NSHomeDirectory()) -> [String] {
        var dirs = [
            "/opt/homebrew/bin", "/opt/homebrew/sbin",
            "/usr/local/bin", "/usr/local/sbin",
            "/opt/local/bin", "/opt/local/sbin",
            "\(home)/.volta/bin",
        ]
        if let nvmDir = nvmDefaultNodeBinDirectory(home: home) {
            dirs.append(nvmDir)
        }
        return dirs
    }

    /// nvm has no fixed "current" symlink usable in PATH, but it does write its default version as
    /// plain text to `~/.nvm/alias/default` — reading that file (no subprocess, no shell) lets us
    /// construct the versioned bin directory nvm's own shell integration would otherwise add.
    nonisolated static func nvmDefaultNodeBinDirectory(home: String, fileManager: FileManager = .default) -> String? {
        let aliasPath = "\(home)/.nvm/alias/default"
        guard let raw = try? String(contentsOfFile: aliasPath, encoding: .utf8) else { return nil }
        let version = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !version.isEmpty else { return nil }
        let versionDir = version.hasPrefix("v") ? version : "v\(version)"
        return "\(home)/.nvm/versions/node/\(versionDir)/bin"
    }

    /// Prepends every candidate directory that exists and isn't already on `basePath`. Pure,
    /// synchronous, and instant — safe to call from anywhere, including MainActor, with no risk of
    /// blocking.
    nonisolated static func augmentedPath(
        basePath: String, home: String = NSHomeDirectory(), fileManager: FileManager = .default
    ) -> String {
        let existing = Set(basePath.split(separator: ":").map(String.init))
        let toAdd = candidateDirectories(home: home).filter { !existing.contains($0) && fileManager.fileExists(atPath: $0) }
        guard !toAdd.isEmpty else { return basePath }
        return (toAdd + [basePath]).joined(separator: ":")
    }

    /// Searches `basePath` (already augmented by the caller, typically) for an executable file
    /// named `name`. Returns its full path if found. Pure filesystem lookup — no subprocess.
    nonisolated static func findExecutable(named name: String, in searchPath: String, fileManager: FileManager = .default) -> String? {
        for dir in searchPath.split(separator: ":") {
            let candidate = "\(dir)/\(name)"
            if fileManager.isExecutableFile(atPath: candidate) { return candidate }
        }
        return nil
    }

    /// Homebrew's install prefix, if Homebrew itself is present — used to offer a one-command
    /// "install Node.js" suggestion when a stdio server's command (typically `npx`) can't be found.
    nonisolated static func homebrewPrefix(fileManager: FileManager = .default) -> String? {
        if fileManager.fileExists(atPath: "/opt/homebrew/bin/brew") { return "/opt/homebrew" }
        if fileManager.fileExists(atPath: "/usr/local/bin/brew") { return "/usr/local" }
        return nil
    }
}

// MARK: - Stdio Transport

/// Launches the server as a subprocess and speaks newline-delimited JSON-RPC over its
/// stdin/stdout, exactly as the original (pre-transport-split) `ExternalMCPClient` did.
@MainActor
final class StdioMCPTransport: MCPTransport {
    private let server: ExternalMCPServer
    /// Set by the owner after construction (not an init parameter) — the owner (`ExternalMCPClient`)
    /// needs to capture `self` weakly to wire this up, which Swift only allows once `self` is
    /// fully initialized, i.e. after its own `init` has finished assigning all stored properties.
    var onTerminated: (() -> Void)?

    private var process: Process?
    private var stdinHandle: FileHandle?
    private var readTask: Task<Void, Never>?
    private var stderrTask: Task<Void, Never>?
    private var pendingCalls: [Int: CheckedContinuation<Data, Error>] = [:]
    private var lineBuffer = Data()

    init(server: ExternalMCPServer) {
        self.server = server
    }

    func prepare() async throws {
        // Inherited environment first, with PATH augmented to include common Homebrew/MacPorts/
        // Volta/nvm install locations (see LoginShellEnvironment — launchd's default PATH doesn't
        // include them, so a bare `npx` spawn fails even though it works fine in Terminal), then
        // layer the user-configured vars over it so they can override.
        var environment = ProcessInfo.processInfo.environment
        let basePath = environment["PATH"] ?? "/usr/bin:/bin:/usr/sbin:/sbin"
        let searchPath = LoginShellEnvironment.augmentedPath(basePath: basePath)
        environment["PATH"] = searchPath
        for (key, value) in server.env { environment[key] = value }

        let proc = Process()
        if server.command.hasPrefix("/") {
            guard FileManager.default.isExecutableFile(atPath: server.command) else {
                throw MCPClientError.commandNotFound(server.command)
            }
            proc.executableURL = URL(fileURLWithPath: server.command)
            proc.arguments = server.args
        } else {
            // Check up front rather than letting `/usr/bin/env <cmd>` fail asynchronously after
            // launch: a missing command is a permanent condition — surfacing it immediately as a
            // clear error, instead of after 3 rounds of crash/restart backoff, is the whole point.
            guard LoginShellEnvironment.findExecutable(named: server.command, in: searchPath) != nil else {
                throw MCPClientError.commandNotFound(server.command)
            }
            proc.executableURL = URL(fileURLWithPath: "/usr/bin/env")
            proc.arguments = [server.command] + server.args
        }
        proc.environment = environment

        let stdinPipe = Pipe()
        let stdoutPipe = Pipe()
        let stderrPipe = Pipe()
        proc.standardInput = stdinPipe
        proc.standardOutput = stdoutPipe
        proc.standardError = stderrPipe

        proc.terminationHandler = { [weak self] _ in
            Task { @MainActor [weak self] in self?.handleProcessTerminated() }
        }

        do {
            try proc.run()
        } catch {
            throw MCPClientError.processLaunchFailed(error.localizedDescription)
        }

        process = proc
        stdinHandle = stdinPipe.fileHandleForWriting
        startReadLoop(pipe: stdoutPipe)
        startStderrLoop(pipe: stderrPipe)
    }

    func sendRequest(_ message: [String: Any], id: Int, timeoutSeconds: Double) async throws -> Data {
        try writeJSON(message)

        // Schedule timeout: cancels the specific pending call by ID
        let timeoutId = id
        Task { [weak self, timeoutId] in
            try? await Task.sleep(nanoseconds: UInt64(timeoutSeconds * 1_000_000_000))
            self?.cancelPendingCall(id: timeoutId, with: MCPClientError.timeout)
        }

        return try await withCheckedThrowingContinuation { cont in
            pendingCalls[id] = cont
        }
    }

    func sendNotification(_ message: [String: Any]) async throws {
        try writeJSON(message)
    }

    func stop() {
        readTask?.cancel()
        stderrTask?.cancel()
        process?.terminate()
        process = nil
        stdinHandle = nil
        lineBuffer = Data()
        for (_, cont) in pendingCalls { cont.resume(throwing: MCPClientError.notConnected) }
        pendingCalls.removeAll()
    }

    // MARK: I/O Loops (detached from MainActor)

    private func startReadLoop(pipe: Pipe) {
        readTask = Task.detached { [weak self] in
            let handle = pipe.fileHandleForReading
            while true {
                let data = handle.availableData
                if data.isEmpty { break }
                await self?.receiveData(data)
            }
        }
    }

    private func startStderrLoop(pipe: Pipe) {
        let name = server.name
        stderrTask = Task.detached {
            let handle = pipe.fileHandleForReading
            var buf = Data()
            while true {
                let data = handle.availableData
                if data.isEmpty { break }
                buf.append(data)
                while let idx = buf.firstIndex(of: UInt8(ascii: "\n")) {
                    let line = String(data: buf[buf.startIndex..<idx], encoding: .utf8) ?? ""
                    buf = Data(buf[buf.index(after: idx)...])
                    if !line.trimmingCharacters(in: .whitespaces).isEmpty {
                        Log.extMcp.warning("[\(name)] \(line)")
                    }
                }
            }
        }
    }

    // MARK: Data Processing (MainActor)

    private func receiveData(_ data: Data) {
        lineBuffer.append(data)
        while let idx = lineBuffer.firstIndex(of: UInt8(ascii: "\n")) {
            let lineData = Data(lineBuffer[lineBuffer.startIndex..<idx])
            lineBuffer = Data(lineBuffer[lineBuffer.index(after: idx)...])
            processLine(lineData)
        }
    }

    private func processLine(_ data: Data) {
        guard !data.isEmpty,
              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let id = json["id"] as? Int,
              let cont = pendingCalls.removeValue(forKey: id) else { return }

        do {
            let resultData = try MCPTransportSupport.extractResult(from: json)
            cont.resume(returning: resultData)
        } catch {
            cont.resume(throwing: error)
        }
    }

    private func cancelPendingCall(id: Int, with error: Error) {
        pendingCalls.removeValue(forKey: id)?.resume(throwing: error)
    }

    private func writeJSON(_ message: [String: Any]) throws {
        guard let handle = stdinHandle, process?.isRunning == true else {
            throw MCPClientError.writeFailed
        }
        guard let data = try? JSONSerialization.data(withJSONObject: message),
              let line = String(data: data, encoding: .utf8) else {
            throw MCPClientError.writeFailed
        }
        do {
            try handle.write(contentsOf: Data((line + "\n").utf8))
        } catch {
            throw MCPClientError.writeFailed
        }
    }

    private func handleProcessTerminated() {
        for (_, cont) in pendingCalls { cont.resume(throwing: MCPClientError.notConnected) }
        pendingCalls.removeAll()
        onTerminated?()
    }
}

// MARK: - HTTP (Streamable HTTP) Transport

/// Speaks MCP's Streamable HTTP transport (spec revision 2025-06-18): every JSON-RPC message is
/// its own HTTP POST to the server's single MCP endpoint. The server may answer with a plain
/// `application/json` body, or open a `text/event-stream` (SSE) response — this transport
/// supports both, but (since Confab only needs request/response tool calls, not server-initiated
/// push) does not open a standalone listening GET stream for unsolicited server messages.
@MainActor
final class HTTPMCPTransport: MCPTransport {
    private let server: ExternalMCPServer
    private let urlSession: URLSession
    private var sessionId: String?

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

    func prepare() async throws {
        guard URL(string: server.url) != nil else {
            throw MCPClientError.invalidConfiguration("Invalid server URL: \(server.url)")
        }
        // No persistent connection to establish ahead of time — the first request (`initialize`)
        // both opens the session and confirms the server is reachable.
    }

    func sendRequest(_ message: [String: Any], id: Int, timeoutSeconds: Double) async throws -> Data {
        let request = try buildRequest(for: message, timeoutSeconds: timeoutSeconds)
        let (data, response) = try await urlSession.data(for: request)
        guard let http = response as? HTTPURLResponse else {
            throw MCPClientError.invalidResponse("Non-HTTP response")
        }
        if let newSessionId = http.value(forHTTPHeaderField: "Mcp-Session-Id") {
            sessionId = newSessionId
        }
        guard (200...299).contains(http.statusCode) else {
            throw MCPClientError.invalidResponse("HTTP \(http.statusCode)")
        }
        let json = try Self.parseResponseBody(
            data, contentType: http.value(forHTTPHeaderField: "Content-Type"), expectedId: id
        )
        return try MCPTransportSupport.extractResult(from: json)
    }

    func sendNotification(_ message: [String: Any]) async throws {
        // Notifications carry no "id", so per spec the server responds 202 Accepted, no body.
        let request = try buildRequest(for: message, timeoutSeconds: 15)
        let (_, response) = try await urlSession.data(for: request)
        guard let http = response as? HTTPURLResponse, (200...299).contains(http.statusCode) else {
            throw MCPClientError.invalidResponse("Notification not accepted")
        }
    }

    func stop() {
        sessionId = nil
    }

    private func buildRequest(for message: [String: Any], timeoutSeconds: Double) throws -> URLRequest {
        guard let url = URL(string: server.url) else {
            throw MCPClientError.invalidConfiguration("Invalid server URL: \(server.url)")
        }
        var request = URLRequest(url: url, timeoutInterval: timeoutSeconds)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.setValue("application/json, text/event-stream", forHTTPHeaderField: "Accept")
        request.setValue("2025-06-18", forHTTPHeaderField: "MCP-Protocol-Version")
        if !server.bearerToken.isEmpty {
            request.setValue("Bearer \(server.bearerToken)", forHTTPHeaderField: "Authorization")
        }
        for (key, value) in server.headers {
            request.setValue(value, forHTTPHeaderField: key)
        }
        if let sessionId {
            request.setValue(sessionId, forHTTPHeaderField: "Mcp-Session-Id")
        }
        request.httpBody = try JSONSerialization.data(withJSONObject: message)
        return request
    }

    /// Parses either a direct `application/json` body, or an SSE (`text/event-stream`) body —
    /// scanning its `data:` lines for the JSON-RPC message whose `id` matches `expectedId` (the
    /// server may send unrelated requests/notifications on the same stream first, per spec).
    nonisolated static func parseResponseBody(_ data: Data, contentType: String?, expectedId: Int) throws -> [String: Any] {
        if contentType?.contains("text/event-stream") == true {
            guard let text = String(data: data, encoding: .utf8) else {
                throw MCPClientError.invalidResponse("Non-UTF8 SSE body")
            }
            for line in text.components(separatedBy: "\n") {
                let trimmed = line.trimmingCharacters(in: .whitespaces)
                guard trimmed.hasPrefix("data:") else { continue }
                let payload = trimmed.dropFirst("data:".count).trimmingCharacters(in: .whitespaces)
                guard let payloadData = payload.data(using: .utf8),
                      let json = try? JSONSerialization.jsonObject(with: payloadData) as? [String: Any],
                      let id = json["id"] as? Int, id == expectedId
                else { continue }
                return json
            }
            throw MCPClientError.invalidResponse("No matching response in SSE stream")
        }
        guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw MCPClientError.invalidResponse("Malformed JSON response body")
        }
        return json
    }
}