Wakeful/Sources/Wakeful/PowerMonitor.swift

import AppKit
import IOKit.ps

/// Watches the two things the smart-activation rules care about: power source
/// (AC vs battery, plus charge level) and whether an external display is attached.
final class PowerMonitor {

    /// Called on the main thread whenever power state or screen layout changes.
    var onChange: (() -> Void)?

    private(set) var isOnAC: Bool = true
    private(set) var batteryPercent: Int?

    private var runLoopSource: CFRunLoopSource?
    private var screenObserver: NSObjectProtocol?

    /// Desktops and Mac minis have no battery; treat them as permanently on AC.
    private(set) var hasBattery: Bool = false

    /// True when any display other than the built-in panel is connected.
    ///
    /// Deliberately not `NSScreen.screens.count > 1`, which is wrong exactly
    /// where the rule matters: in clamshell mode the built-in panel leaves
    /// `NSScreen.screens`, a mirror set is a single NSScreen, and a desktop
    /// Mac's only monitor is external. The online list has every connected
    /// display, hardware mirrors included. DisplayLink, Sidecar and AirPlay
    /// displays are not built in, so they count too.
    var externalDisplayConnected: Bool {
        var count: UInt32 = 0
        guard CGGetOnlineDisplayList(0, nil, &count) == .success, count > 0 else {
            return NSScreen.screens.count > 1
        }
        var displays = [CGDirectDisplayID](repeating: 0, count: Int(count))
        guard CGGetOnlineDisplayList(count, &displays, &count) == .success else {
            return NSScreen.screens.count > 1
        }
        return displays.prefix(Int(count)).contains { CGDisplayIsBuiltin($0) == 0 }
    }

    func start() {
        guard runLoopSource == nil else { return }
        refresh()

        let context = Unmanaged.passUnretained(self).toOpaque()
        let callback: IOPowerSourceCallbackType = { rawContext in
            guard let rawContext else { return }
            let monitor = Unmanaged<PowerMonitor>.fromOpaque(rawContext).takeUnretainedValue()
            monitor.refresh()
            monitor.onChange?()
        }

        if let source = IOPSNotificationCreateRunLoopSource(callback, context)?.takeRetainedValue() {
            runLoopSource = source
            // .commonModes, not .defaultMode: otherwise power changes are dropped
            // while a menu is tracking or an alert is up.
            CFRunLoopAddSource(CFRunLoopGetMain(), source, CFRunLoopMode.commonModes)
        }

        screenObserver = NotificationCenter.default.addObserver(
            forName: NSApplication.didChangeScreenParametersNotification,
            object: nil,
            queue: .main
        ) { [weak self] _ in
            self?.onChange?()
        }
    }

    func stop() {
        if let runLoopSource {
            CFRunLoopRemoveSource(CFRunLoopGetMain(), runLoopSource, CFRunLoopMode.commonModes)
        }
        runLoopSource = nil
        if let screenObserver {
            NotificationCenter.default.removeObserver(screenObserver)
        }
        screenObserver = nil
    }

    /// Reads a fresh power-source snapshot.
    ///
    /// AC state is derived from the internal battery's own power-source state
    /// rather than `IOPSGetProvidingPowerSourceType`, which keeps the Unmanaged
    /// ownership rules unambiguous.
    func refresh() {
        guard let snapshot = IOPSCopyPowerSourcesInfo()?.takeRetainedValue(),
              let list = IOPSCopyPowerSourcesList(snapshot)?.takeRetainedValue() as NSArray?
        else {
            hasBattery = false
            isOnAC = true
            batteryPercent = nil
            return
        }

        var foundBattery = false
        var onAC = true
        var percent: Int?

        for element in list {
            let source = element as CFTypeRef
            guard let description =
                    IOPSGetPowerSourceDescription(snapshot, source)?.takeUnretainedValue()
                    as? [String: Any]
            else { continue }

            guard let type = description[kIOPSTypeKey] as? String,
                  type == kIOPSInternalBatteryType else { continue }

            foundBattery = true

            if let state = description[kIOPSPowerSourceStateKey] as? String {
                onAC = (state == kIOPSACPowerValue)
            }

            if let current = description[kIOPSCurrentCapacityKey] as? Int,
               let maximum = description[kIOPSMaxCapacityKey] as? Int,
               maximum > 0 {
                percent = Int((Double(current) / Double(maximum) * 100.0).rounded())
            }
            break
        }

        hasBattery = foundBattery
        isOnAC = foundBattery ? onAC : true
        batteryPercent = percent
    }

    deinit {
        stop()
    }
}