nucleic/Sources/NucleicApp/PowerHelperClient.swift

import Foundation
import NucleicPowerProtocol
import ServiceManagement

enum PowerHelperState: Equatable {
    case notRegistered
    case requiresApproval
    case enabled
    case notFound
    case failed(String)
}

enum PowerHelperError: LocalizedError {
    case unavailable
    case connection(String)
    case remote(Error)

    var errorDescription: String? {
        switch self {
        case .unavailable:
            "In System Settings, go to General > Login Items & Extensions, then turn on "
                + "Nucleic under Allow in the Background."
        case .connection(let detail):
            "Could not connect to Nucleic's power helper: \(detail)"
        case .remote(let error):
            error.localizedDescription
        }
    }
}

private final class XPCReplyGate: @unchecked Sendable {
    private let lock = NSLock()
    private var resumed = false

    func once(_ body: @Sendable () -> Void) {
        lock.lock()
        defer { lock.unlock() }
        guard !resumed else { return }
        resumed = true
        body()
    }
}

@MainActor
final class PowerHelperClient {
    static let shared = PowerHelperClient()

    private var connection: NSXPCConnection?
    var connectionLost: (() -> Void)?

    private init() {}

    /// The daemon's current registration status, queried **off the main actor**. `SMAppService.status`
    /// is a synchronous XPC round-trip to `smd` that can take hundreds of ms on a loaded or beta system,
    /// so it must never run on the main thread. A fresh `SMAppService.daemon` handle is built inside the
    /// detached task: the handle is a lightweight reference keyed by plist name — status and register act
    /// on `smd`'s own state, not on instance state — so this is equivalent to querying a stored `service`
    /// without pinning the blocking call to the main actor.
    nonisolated static func currentState() async -> PowerHelperState {
        await Task.detached(priority: .utility) {
            Self.map(SMAppService.daemon(plistName: NucleicPowerHelperConstants.plistName).status)
        }.value
    }

    /// Register the daemon if it isn't yet, returning the resulting state. Runs entirely off the main
    /// actor — both the status check and `register()` are synchronous `smd` XPC.
    @discardableResult
    nonisolated static func registerIfNeeded() async -> PowerHelperState {
        await Task.detached(priority: .utility) {
            let service = SMAppService.daemon(plistName: NucleicPowerHelperConstants.plistName)
            // Already usable, or already waiting on the user — nothing left to (re)register.
            switch service.status {
            case .enabled, .requiresApproval:
                return Self.map(service.status)
            default:
                break
            }
            // Everything else — `.notRegistered` AND `.notFound` — is a pre-registration state, so
            // attempt `register()`. This is load-bearing: for a *bundled* daemon that macOS hasn't
            // registered yet, `status` frequently reports `.notFound` (not `.notRegistered`) until the
            // first `register()` call introduces the bundled plist to BTM/LaunchServices — typically on
            // the first launch after install. Gating registration on `status == .notRegistered` there
            // skips `register()` forever: no approval prompt, no Login Items entry, no error, and no
            // `smd` activity at all (the exact "Smart Sleep never registers" failure). Always call
            // `register()` and let it either advance the service to `.requiresApproval` (surfacing the
            // prompt) or throw a concrete `.failed` we can show — never silently no-op on `.notFound`.
            do {
                try service.register()
                return Self.map(service.status)
            } catch {
                return .failed(error.localizedDescription)
            }
        }.value
    }

    private nonisolated static func map(_ status: SMAppService.Status) -> PowerHelperState {
        switch status {
        case .notRegistered: .notRegistered
        case .requiresApproval: .requiresApproval
        case .enabled: .enabled
        case .notFound: .notFound
        @unknown default: .notFound
        }
    }

    func openApprovalSettings() {
        SMAppService.openSystemSettingsLoginItems()
    }

    /// Acquire (`active`) or release the privileged full-wake lease over XPC. The caller must have
    /// confirmed the helper is `.enabled` (via ``currentState()`` / ``registerIfNeeded()``) beforehand:
    /// this performs **no** status query of its own, so the reconcile loop makes exactly one synchronous
    /// `smd` round-trip per iteration instead of three. A connect to a non-enabled service simply fails
    /// through the XPC error handler below.
    func setFullWakeActive(_ active: Bool) async throws {
        let connection = connection ?? makeConnection()
        self.connection = connection

        try await withCheckedThrowingContinuation {
            (continuation: CheckedContinuation<Void, Error>) in
            let gate = XPCReplyGate()
            let proxy = connection.remoteObjectProxyWithErrorHandler { error in
                gate.once { continuation.resume(throwing: PowerHelperError.connection(
                    error.localizedDescription)) }
            }
            guard let helper = proxy as? NucleicPowerHelperProtocol else {
                gate.once { continuation.resume(throwing:
                    PowerHelperError.connection("invalid XPC proxy")) }
                return
            }
            helper.setFullWakeActive(active) { error in
                gate.once {
                    if let error {
                        continuation.resume(throwing: PowerHelperError.remote(error))
                    } else {
                        continuation.resume()
                    }
                }
            }
        }
    }

    private func makeConnection() -> NSXPCConnection {
        let connection = NSXPCConnection(
            machServiceName: NucleicPowerHelperConstants.serviceName,
            options: .privileged)
        connection.remoteObjectInterface = NSXPCInterface(
            with: NucleicPowerHelperProtocol.self)
        connection.invalidationHandler = { [weak self] in
            Task { @MainActor in
                self?.connection = nil
                self?.connectionLost?()
            }
        }
        connection.interruptionHandler = { [weak self] in
            Task { @MainActor in
                self?.connection = nil
                self?.connectionLost?()
            }
        }
        connection.resume()
        return connection
    }
}