import Foundation
import ScreenCaptureKit
import CoreGraphics
import ImageIO

@available(macOS 14.0, *)
func captureScreen(to outputPath: String) async throws {
    let content = try await SCShareableContent.excludingDesktopWindows(
        false, onScreenWindowsOnly: true
    )
    let mainID = CGMainDisplayID()
    guard let display = content.displays.first(where: { $0.displayID == mainID })
        ?? content.displays.first
    else {
        throw NSError(
            domain: "VisionClaw", code: 1,
            userInfo: [NSLocalizedDescriptionKey: "No display found"]
        )
    }

    let filter = SCContentFilter(display: display, excludingWindows: [])
    let config = SCStreamConfiguration()
    config.showsCursor = true
    config.captureResolution = .best
    // Set output dimensions large enough to avoid downscaling on Retina displays.
    // With .best, ScreenCaptureKit captures at native resolution; if these exceed
    // the actual pixel count it won't upscale.
    config.width = display.width * 2
    config.height = display.height * 2

    let image = try await SCScreenshotManager.captureImage(
        contentFilter: filter, configuration: config
    )
    try writePNG(image, to: outputPath)
}

func writePNG(_ image: CGImage, to path: String) throws {
    let url = URL(fileURLWithPath: path) as CFURL
    guard let dest = CGImageDestinationCreateWithURL(
        url, "public.png" as CFString, 1, nil
    ) else {
        throw NSError(
            domain: "VisionClaw", code: 2,
            userInfo: [NSLocalizedDescriptionKey: "Cannot create image destination"]
        )
    }
    CGImageDestinationAddImage(dest, image, nil)
    guard CGImageDestinationFinalize(dest) else {
        throw NSError(
            domain: "VisionClaw", code: 3,
            userInfo: [NSLocalizedDescriptionKey: "Cannot finalize PNG"]
        )
    }
}

let outputPath = CommandLine.arguments.count > 1
    ? CommandLine.arguments[1]
    : "/tmp/screenshot.png"

// Hard timeout guard: ScreenCaptureKit's async APIs (SCShareableContent,
// SCScreenshotManager.captureImage) are known to hang indefinitely under
// certain system conditions — e.g. right after a Screen Recording permission
// change, during display sleep/wake transitions, external display hot-plug,
// user switching, or a stuck WindowServer. When that happens the awaiting
// Task never resumes, dispatchMain() never returns, and the helper process
// sits forever consuming CPU. The Node parent (execAsync) also has no
// default timeout, so the whole agent main loop deadlocks on the screenshot
// call.
//
// This guard force-exits the process after a deadline so the parent always
// observes a non-zero exit code and can fall back to `screencapture -x -C`.
// SIGKILL from the parent still works if this timer itself gets starved,
// but this is the primary defence.
let timeoutSec: Int = {
    if let raw = ProcessInfo.processInfo.environment["VISIONCLAW_SCREENSHOT_TIMEOUT_SEC"],
       let n = Int(raw), n > 0 {
        return n
    }
    return 10
}()
DispatchQueue.global().asyncAfter(deadline: .now() + .seconds(timeoutSec)) {
    fputs("Error: screenshot capture timed out after \(timeoutSec)s\n", stderr)
    Foundation.exit(124)
}

if #available(macOS 14.0, *) {
    Task {
        do {
            try await captureScreen(to: outputPath)
            Foundation.exit(0)
        } catch {
            fputs("Error: \(error.localizedDescription)\n", stderr)
            Foundation.exit(1)
        }
    }
    dispatchMain()
} else {
    fputs("Requires macOS 14.0+\n", stderr)
    Foundation.exit(1)
}
