kapture/Sources/Kapture/CLIService.swift

import AppKit
import ImageIO
import UniformTypeIdentifiers

/// Serves `kapture` CLI requests from the running menu-bar app. The CLI
/// writes a JSON request into `~/Library/Application Support/Kapture/cli`
/// and posts a distributed notification carrying its id; this service
/// captures and writes the JSON response alongside it. Captures use the
/// app's own Screen Recording grant, so the CLI needs no permissions.
final class CLIService {
    static let shared = CLIService()

    static let notificationName = "dev.tzgyn.kapture.cli.request"
    static let requestKey = "id"

    private var observer: NSObjectProtocol?

    private var directory: URL {
        FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
            .appendingPathComponent("Kapture/cli", isDirectory: true)
    }

    private func requestURL(_ id: String) -> URL {
        directory.appendingPathComponent("\(id).request.json")
    }

    private func responseURL(_ id: String) -> URL {
        directory.appendingPathComponent("\(id).response.json")
    }

    func start() {
        guard observer == nil else { return }
        try? FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
        observer = DistributedNotificationCenter.default().addObserver(
            forName: .init(Self.notificationName), object: nil, queue: .main
        ) { [weak self] note in
            guard let id = note.userInfo?[Self.requestKey] as? String else { return }
            self?.handle(id: id)
        }
        CaptureController.log("cli: listening")
    }

    // MARK: - Requests

    private func handle(id: String) {
        let url = requestURL(id)
        defer { try? FileManager.default.removeItem(at: url) }
        guard let data = try? Data(contentsOf: url),
              let request = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any] else {
            respond(id: id, ["ok": false, "error": "Malformed request"])
            return
        }
        switch request["command"] as? String {
        case "list":
            respond(id: id, [
                "ok": true,
                "windows": windows().map { window in
                    [
                        "id": Int(window.id),
                        "app": window.app,
                        "title": window.title,
                        "width": Int(window.bounds.width),
                        "height": Int(window.bounds.height),
                    ] as [String: Any]
                },
            ])
        case "capture":
            handleCapture(id: id, request: request)
        default:
            respond(id: id, ["ok": false, "error": "Unknown command"])
        }
    }

    private func handleCapture(id: String, request: [String: Any]) {
        guard CGPreflightScreenCaptureAccess() else {
            respond(id: id, ["ok": false, "error": "Kapture does not have Screen Recording permission"])
            return
        }
        let outPath = request["out"] as? String

        if request["screen"] as? Bool == true {
            guard let screen = NSScreen.screens.first else {
                respond(id: id, ["ok": false, "error": "No screen available"])
                return
            }
            capture(rect: screen.frame, windowID: nil, out: outPath, id: id)
            return
        }

        if let number = request["windowID"] as? NSNumber,
           let window = windows().first(where: { Int($0.id) == number.intValue }) {
            capture(rect: window.bounds, windowID: window.id, out: outPath, id: id)
            return
        }

        guard let window = bestWindow(for: request["app"] as? String) else {
            let error = (request["app"] as? String).map { "No on-screen window found for app \"\($0)\"" }
                ?? "No window found"
            respond(id: id, ["ok": false, "error": error])
            return
        }
        capture(rect: window.bounds, windowID: window.id, out: outPath, id: id)
    }

    private func capture(rect: NSRect, windowID: CGWindowID?, out: String?, id: String) {
        ScreenCapture.shared.capture(rect, excludingWindowIDs: [], windowID: windowID) { [weak self] image in
            DispatchQueue.main.async {
                guard let self else { return }
                guard let image else {
                    self.respond(id: id, ["ok": false, "error": "Capture failed (Screen Recording permission?)"])
                    return
                }
                let url = out.map { URL(fileURLWithPath: $0) }
                    ?? FileManager.default.temporaryDirectory.appendingPathComponent("kapture-\(Self.timestamp()).png")
                guard let dest = CGImageDestinationCreateWithURL(url as CFURL, UTType.png.identifier as CFString, 1, nil) else {
                    self.respond(id: id, ["ok": false, "error": "Could not create \(url.path)"])
                    return
                }
                CGImageDestinationAddImage(dest, image, nil)
                guard CGImageDestinationFinalize(dest) else {
                    self.respond(id: id, ["ok": false, "error": "Could not write \(url.path)"])
                    return
                }
                CaptureController.log("cli: captured \(image.width)x\(image.height) -> \(url.path)")
                self.respond(id: id, ["ok": true, "path": url.path, "width": image.width, "height": image.height])
            }
        }
    }

    private func respond(id: String, _ body: [String: Any]) {
        guard let data = try? JSONSerialization.data(withJSONObject: body) else { return }
        try? data.write(to: responseURL(id), options: .atomic)
    }

    // MARK: - Windows

    private struct WindowInfo {
        let id: CGWindowID
        let pid: Int
        let app: String
        let title: String
        let bounds: NSRect
    }

    /// On-screen layer-0 windows, front to back, own panels excluded.
    private func windows() -> [WindowInfo] {
        guard let list = CGWindowListCopyWindowInfo(
            [.optionOnScreenOnly, .excludeDesktopElements], kCGNullWindowID
        ) as? [[String: Any]] else { return [] }

        let ownPid = Int(ProcessInfo.processInfo.processIdentifier)
        return list.compactMap { info in
            guard let id = info[kCGWindowNumber as String] as? CGWindowID,
                  let layer = info[kCGWindowLayer as String] as? Int, layer == 0,
                  let pid = info[kCGWindowOwnerPID as String] as? Int, pid != ownPid,
                  let bounds = info[kCGWindowBounds as String] as? [String: CGFloat],
                  (bounds["Width"] ?? 0) >= 80, (bounds["Height"] ?? 0) >= 60 else { return nil }
            let cgRect = NSRect(
                x: bounds["X"] ?? 0, y: bounds["Y"] ?? 0,
                width: bounds["Width"] ?? 0, height: bounds["Height"] ?? 0
            )
            return WindowInfo(
                id: id,
                pid: pid,
                app: info[kCGWindowOwnerName as String] as? String ?? "",
                title: info[kCGWindowName as String] as? String ?? "",
                bounds: Self.cgToAppKit(cgRect)
            )
        }
    }

    private func bestWindow(for query: String?) -> WindowInfo? {
        let candidates = windows()
        guard let query, !query.isEmpty else {
            if let front = NSWorkspace.shared.frontmostApplication,
               let window = candidates.first(where: { $0.pid == front.processIdentifier }) {
                return window
            }
            return candidates.first
        }
        let needle = query.lowercased()
        if let window = candidates.first(where: { $0.app.lowercased().contains(needle) }) {
            return window
        }
        for app in NSWorkspace.shared.runningApplications {
            let name = (app.localizedName ?? "").lowercased()
            let bundle = (app.bundleIdentifier ?? "").lowercased()
            if name.contains(needle) || bundle.contains(needle),
               let window = candidates.first(where: { $0.pid == app.processIdentifier }) {
                return window
            }
        }
        return nil
    }

    private static func cgToAppKit(_ rect: NSRect) -> NSRect {
        guard let main = NSScreen.screens.first else { return rect }
        return NSRect(
            x: rect.origin.x,
            y: main.frame.height - rect.origin.y - rect.height,
            width: rect.width,
            height: rect.height
        )
    }

    private static func timestamp() -> String {
        let formatter = DateFormatter()
        formatter.dateFormat = "yyyyMMdd-HHmmss"
        return formatter.string(from: Date())
    }
}