import AppKit

struct ToolError: Error {
    let message: String

    init(_ message: String) {
        self.message = message
    }
}

enum Content {
    case text(String)
    case png(Data)

    var json: [String: Any] {
        switch self {
        case .text(let s):
            return ["type": "text", "text": s]
        case .png(let d):
            return ["type": "image", "data": d.base64EncodedString(), "mimeType": "image/png"]
        }
    }
}

struct Tool {
    let name: String
    let description: String
    let schema: String
    let run: ([String: Any]) async throws -> [Content]
}

let stdoutLock = NSLock()

let mediaDir: String? = {
    let args = CommandLine.arguments
    guard let i = args.firstIndex(of: "--media-dir"), i + 1 < args.count else {
        return nil
    }
    return args[i + 1]
}()

@main
struct Server {
    static func main() {
        _ = NSApplication.shared
        NSApp.setActivationPolicy(.prohibited)
        AXUIElementSetMessagingTimeout(AXUIElementCreateSystemWide(), 3)
        Thread.detachNewThread {
            let inFlight = DispatchGroup()
            while let line = readLine() {
                inFlight.enter()
                Task.detached {
                    await handle(line)
                    inFlight.leave()
                }
            }
            inFlight.wait()
            exit(0)
        }
        RunLoop.main.run()
    }

    static func handle(_ line: String) async {
        guard let data = line.data(using: .utf8),
              let req = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let id = req["id"]
        else {
            return
        }
        var resp: [String: Any] = ["jsonrpc": "2.0", "id": id]
        let params = req["params"] as? [String: Any] ?? [:]
        switch req["method"] as? String ?? "" {
        case "initialize":
            resp["result"] = [
                "protocolVersion": params["protocolVersion"] as? String ?? "2025-06-18",
                "capabilities": ["tools": [:]],
                "serverInfo": ["name": "hands", "version": "0.0.0"],
            ]
        case "ping":
            resp["result"] = [:] as [String: Any]
        case "tools/list":
            resp["result"] = ["tools": tools.map { t in
                ["name": t.name, "description": t.description, "inputSchema": try! JSONSerialization.jsonObject(with: Data(t.schema.utf8))]
            }]
        case "tools/call":
            resp["result"] = await call(params)
        default:
            resp["error"] = ["code": -32601, "message": "method not found: \(req["method"] ?? "")"]
        }
        var out = try! JSONSerialization.data(withJSONObject: resp, options: .withoutEscapingSlashes)
        out.append(Data("\n".utf8))
        stdoutLock.withLock {
            FileHandle.standardOutput.write(out)
        }
    }

    static func call(_ params: [String: Any]) async -> [String: Any] {
        let name = params["name"] as? String ?? ""
        let args = params["arguments"] as? [String: Any] ?? [:]
        guard let tool = tools.first(where: { $0.name == name }) else {
            return ["content": [Content.text("unknown tool: \(name)").json], "isError": true]
        }
        do {
            let content = try await tool.run(args)
            return ["content": content.map { $0.json }]
        } catch let e as ToolError {
            return ["content": [Content.text(e.message).json], "isError": true]
        } catch {
            return ["content": [Content.text("\(error)").json], "isError": true]
        }
    }
}