import CoreGraphics
import IOKit.pwr_mgt
import ScreenCaptureKit

func displayAsleep() -> Bool {
    let id = CGMainDisplayID()
    return CGDisplayIsAsleep(id) != 0 || CGDisplayBounds(id).width == 0
}

func displayCapturable() async -> Bool {
    if displayAsleep() {
        return false
    }
    // Without Screen Recording this throws; stop waiting so tools that don't capture still run.
    guard let content = try? await SCShareableContent.excludingDesktopWindows(false, onScreenWindowsOnly: true) else {
        return true
    }
    return content.displays.contains(where: { $0.displayID == CGMainDisplayID() })
}

func withDisplayAwake<T>(_ body: () async throws -> T) async throws -> T {
    var hold = IOPMAssertionID(0)
    IOPMAssertionCreateWithName(kIOPMAssertionTypePreventUserIdleDisplaySleep as CFString, IOPMAssertionLevel(kIOPMAssertionLevelOn), "byre-hands" as CFString, &hold)
    defer {
        IOPMAssertionRelease(hold)
    }
    if displayAsleep() {
        var wake = IOPMAssertionID(0)
        IOPMAssertionDeclareUserActivity("byre-hands" as CFString, kIOPMUserActiveLocal, &wake)
        defer {
            IOPMAssertionRelease(wake)
        }
        let deadline = ContinuousClock.now + .seconds(10)
        while !(await displayCapturable()) {
            if ContinuousClock.now > deadline {
                throw ToolError("the display didn't wake within 10s")
            }
            try await Task.sleep(for: .milliseconds(100))
        }
    }
    return try await body()
}