import Foundation
import IOKit.pwr_mgt
import NucleicCore
import Observation

/// Smart Sleep has two power tiers:
/// - Enabled with no local turn: `PreventSystemSleep` lets macOS use Dark Wake, keeping mesh
///   maintenance possible without holding the machine in a full user wake.
/// - A local turn in flight: the root helper leases `pmset disablesleep 1` (the lid-close
///   override) while a normal idle-sleep assertion keeps the CPU fully available to the agent.
///
@MainActor
@Observable
final class SleepBlocker {
    static let shared = SleepBlocker()
    private init() {
        PowerHelperClient.shared.connectionLost = { [weak self] in
            self?.helperConnectionLost()
        }
    }

    private static let preventSystemSleep = "PreventSystemSleep" as CFString
    private static let noAssertion = IOPMAssertionID(0)
    static let enabledKey = "nucleic.power.blockSleep"
    static var isEnabled: Bool {
        UserDefaults.standard.bool(forKey: enabledKey)
    }

    private var darkWakeAssertionID = noAssertion
    private var fullWakeActivityAssertionID = noAssertion
    private var idleSleepActivity: NSObjectProtocol?
    private var localTurnCount = 0
    private var desiredFullWake = false
    private var appliedFullWake = false
    private var helperReconcileTask: Task<Void, Never>?

    private(set) var helperState: PowerHelperState = .notRegistered
    private(set) var lastError: String?

    var requiresHelperApproval: Bool {
        helperState == .requiresApproval || helperState == .notRegistered
    }

    func setEnabled(_ enabled: Bool) {
        if enabled {
            beginDarkWakeAssertion()
        } else {
            endDarkWakeAssertion()
        }
        reconcileLocalTier(enabled: enabled)
        // With a local turn in flight, `reconcileHelperLease` already registers the login item as its
        // first step, so leave it to that. Otherwise (Smart Sleep on but no agent running) the reconcile
        // loop won't touch the helper, so eagerly register here to surface the approval prompt now —
        // off the main actor, since SMAppService.register()/status are synchronous smd XPC that can
        // stall the main thread for hundreds of ms. The `!desiredFullWake` guard keeps the two paths
        // mutually exclusive, so `register()` is never called twice concurrently.
        if enabled, !desiredFullWake {
            Task { [weak self] in
                let state = await PowerHelperClient.registerIfNeeded()
                guard let self else { return }
                self.helperState = state
                if case .failed(let detail) = state {
                    self.lastError = "Could not register Nucleic's background helper: \(detail)"
                }
            }
        }
    }

    func updateLocalTurnCount(_ count: Int) {
        localTurnCount = max(0, count)
        reconcileLocalTier(enabled: Self.isEnabled)
    }

    func applyStoredSetting() {
        setEnabled(Self.isEnabled)
    }

    func openHelperApprovalSettings() {
        PowerHelperClient.shared.openApprovalSettings()
    }

    private func beginDarkWakeAssertion() {
        guard darkWakeAssertionID == Self.noAssertion else { return }
        var assertionID = Self.noAssertion
        let result = IOPMAssertionCreateWithName(
            Self.preventSystemSleep,
            IOPMAssertionLevel(kIOPMAssertionLevelOn),
            "Smart Sleep is on in Nucleic Settings" as CFString,
            &assertionID)
        if result == kIOReturnSuccess {
            darkWakeAssertionID = assertionID
        } else {
            lastError = "Could not create the Dark Wake assertion (IOKit \(result))."
        }
    }

    private func endDarkWakeAssertion() {
        guard darkWakeAssertionID != Self.noAssertion else { return }
        IOPMAssertionRelease(darkWakeAssertionID)
        darkWakeAssertionID = Self.noAssertion
    }

    private func reconcileLocalTier(enabled: Bool) {
        desiredFullWake = enabled && localTurnCount > 0
        if desiredFullWake, idleSleepActivity == nil {
            idleSleepActivity = ProcessInfo.processInfo.beginActivity(
                options: [.idleSystemSleepDisabled, .suddenTerminationDisabled],
                reason: "A local Nucleic agent is running")
        } else if !desiredFullWake, let activity = idleSleepActivity {
            ProcessInfo.processInfo.endActivity(activity)
            idleSleepActivity = nil
            endFullWakeActivity()
        }
        reconcileHelperLease()
    }

    private func reconcileHelperLease() {
        guard helperReconcileTask == nil else { return }
        helperReconcileTask = Task { [weak self] in
            guard let self else { return }
            // The last target we actually processed. Used at the tail to distinguish "desired changed
            // while we were reconciling" (a suppressed trigger to service) from "we stopped because the
            // helper isn't usable" (must NOT rerun, or it spins) — we only rerun when it no longer
            // matches the live desire.
            var lastTarget = self.desiredFullWake
            while self.appliedFullWake != self.desiredFullWake {
                let target = self.desiredFullWake
                lastTarget = target
                // One status query per iteration, off the main actor (registers if not yet registered).
                let state = await PowerHelperClient.registerIfNeeded()
                self.helperState = state
                guard state == .enabled else {
                    // The privileged helper isn't usable: the user hasn't approved the login item
                    // (.requiresApproval / .notRegistered), it isn't installed (.notFound), or the
                    // registration failed. None of these change until the *user* acts, so retrying can
                    // only re-stall the main thread on smd — stop and surface the state in Settings. A
                    // real trigger (Smart Sleep toggled, a turn starting/ending, the app becoming
                    // active) re-enters here via `reconcileHelperLease` / `recheckHelperApproval`.
                    if target {
                        if case .failed(let detail) = state {
                            self.lastError = "Could not register Nucleic's background helper: \(detail)"
                        } else {
                            self.lastError = PowerHelperError.unavailable.errorDescription
                        }
                    } else {
                        // We wanted to release; an unusable helper holds no lease, so we're already done.
                        self.appliedFullWake = false
                        self.endFullWakeActivity()
                    }
                    break
                }
                do {
                    try await PowerHelperClient.shared.setFullWakeActive(target)
                    self.appliedFullWake = target
                    self.lastError = nil
                    if target { self.declareRemoteUserActivity() }
                } catch {
                    self.lastError = error.localizedDescription
                    if !target {
                        // A dead XPC connection releases its server-side lease automatically.
                        self.appliedFullWake = false
                        self.endFullWakeActivity()
                        break
                    }
                    // The helper is enabled but the call failed (a transient XPC hiccup); back off before
                    // retrying rather than spinning.
                    try? await Task.sleep(for: .seconds(5))
                }
            }
            self.helperReconcileTask = nil
            // Rerun only if the desired state changed out from under the pass we just finished (a trigger
            // the `helperReconcileTask` guard suppressed). A pass that stopped on an unusable helper left
            // `lastTarget == desiredFullWake`, so it stays put until a real trigger fires.
            if self.desiredFullWake != lastTarget {
                self.reconcileHelperLease()
            }
        }
    }

    /// Re-evaluate the helper lease on a real trigger (the app becoming active, e.g. the user returning
    /// from approving the login item in System Settings). The reconcile loop treats an unapproved or
    /// missing helper as terminal and stops, so without a nudge like this an approval granted while it
    /// was blocked wouldn't take effect until the next turn boundary. Cheap no-op unless full wake is
    /// wanted but not yet applied.
    func recheckHelperApproval() {
        guard desiredFullWake, !appliedFullWake else { return }
        reconcileHelperLease()
    }

    private func helperConnectionLost() {
        guard appliedFullWake else { return }
        appliedFullWake = false
        endFullWakeActivity()
        reconcileHelperLease()
    }

    /// Promote an already-running Dark Wake after the privileged lid-sleep override is active.
    /// Remote activity keeps the components needed by an agent usable without pretending that a
    /// person touched the local keyboard (which would unnecessarily light an open display).
    private func declareRemoteUserActivity() {
        var assertionID = fullWakeActivityAssertionID
        let result = IOPMAssertionDeclareUserActivity(
            "A local Nucleic agent started during Dark Wake" as CFString,
            kIOPMUserActiveRemote,
            &assertionID)
        if result == kIOReturnSuccess {
            fullWakeActivityAssertionID = assertionID
        } else {
            lastError = "Could not promote Dark Wake for the local agent (IOKit \(result))."
        }
    }

    private func endFullWakeActivity() {
        guard fullWakeActivityAssertionID != Self.noAssertion else { return }
        IOPMAssertionRelease(fullWakeActivityAssertionID)
        fullWakeActivityAssertionID = Self.noAssertion
    }
}
