One Runtime Session, Many Claims

import Foundation

/// Who needs the watch app running with the wrist down. watchOS runs ONE extended runtime
/// session per app, so the parts never start one; they claim it and release it.
struct RuntimeClaims: Equatable {
    enum Claim: Hashable, Sendable { case turn, wait, reading }

    private(set) var claims: Set<Claim> = []
    var wanted: Bool { !claims.isEmpty }

    mutating func claim(_ claim: Claim) { claims.insert(claim) }
    mutating func release(_ claim: Claim) { claims.remove(claim) }

    enum Step: Equatable { case start, keep, endSoon, waitForFront, nothing }

    /// `alive`: a session this app started that the system has not ended (starting or running).
    func step(alive: Bool, appActive: Bool) -> Step {
        switch (wanted, alive) {
        case (true, true): return .keep
        case (true, false): return appActive ? .start : .waitForFront   // a session starts only in front
        case (false, true): return .endSoon
        case (false, false): return .nothing
        }
    }

    /// Kept this long after the last release: the next claim usually follows at once, and an
    /// end followed by a start in the same breath is refused with sessionInProgress.
    static let grace: TimeInterval = 5

    enum End: Equatable { case leftTheApp, expired, anotherSession, other }

    /// Only the refusal for an earlier session of ours is worth a quick second try.
    static func retries(after end: End) -> Bool { end == .anotherSession }
    static let retryDelay: TimeInterval = 1
    static let retrySpacing: TimeInterval = 10
}

import WatchKit

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

    private var claims = RuntimeClaims()
    private var session: WKExtendedRuntimeSession?
    private var ending: Task<Void, Never>?
    private var lastRetry: Date?
    private let delegate = Delegate()                 // retained: the session's delegate is weak

    /// The conversation decides whether it ends with the session (it does for leftTheApp and expired).
    var onSystemEnd: (RuntimeClaims.End) -> Void = { _ in }

    private init() {
        delegate.onEnd = { [weak self] id, end in self?.systemEnded(id, end) }
    }

    func claim(_ c: RuntimeClaims.Claim) { claims.claim(c); apply() }
    func release(_ c: RuntimeClaims.Claim) { claims.release(c); apply() }
    func cameToFront() { apply() }                     // a claim made behind the face starts now

    private var alive: Bool { session.map { $0.state != .invalid } ?? false }

    private func apply() {
        let active = WKApplication.shared().applicationState == .active
        switch claims.step(alive: alive, appActive: active) {
        case .start:
            ending?.cancel(); ending = nil
            let s = WKExtendedRuntimeSession()        // the type comes from WKBackgroundModes
            s.delegate = delegate
            s.start()
            session = s
        case .keep:
            ending?.cancel(); ending = nil            // claimed again inside the grace
        case .endSoon:
            guard ending == nil else { return }
            ending = Task { [weak self] in
                try? await Task.sleep(for: .seconds(RuntimeClaims.grace))
                guard let self, !Task.isCancelled, !self.claims.wanted else { return }
                self.ending = nil
                self.session?.invalidate()
                self.session = nil
            }
        case .waitForFront, .nothing:
            break
        }
    }

    private func systemEnded(_ id: ObjectIdentifier, _ end: RuntimeClaims.End) {
        guard let session, ObjectIdentifier(session) == id else { return }   // an older session's word
        self.session = nil
        guard claims.wanted else { return }
        onSystemEnd(end)
        guard RuntimeClaims.retries(after: end) else { return }
        // Once, a second later, and never more often than every ten seconds: a refusal that keeps
        // coming back waits for the app's next return to the front instead of looping.
        if let lastRetry, Date().timeIntervalSince(lastRetry) < RuntimeClaims.retrySpacing { return }
        lastRetry = Date()
        Task { [weak self] in
            try? await Task.sleep(for: .seconds(RuntimeClaims.retryDelay))
            self?.apply()
        }
    }

    private final class Delegate: NSObject, WKExtendedRuntimeSessionDelegate {
        var onEnd: (@MainActor (ObjectIdentifier, RuntimeClaims.End) -> Void)?

        func extendedRuntimeSessionDidStart(_ s: WKExtendedRuntimeSession) {}
        func extendedRuntimeSessionWillExpire(_ s: WKExtendedRuntimeSession) {}
        func extendedRuntimeSession(_ s: WKExtendedRuntimeSession,
                                    didInvalidateWith reason: WKExtendedRuntimeSessionInvalidationReason,
                                    error: Error?) {
            let end: RuntimeClaims.End = switch reason {
            case .resignedFrontmost: .leftTheApp
            case .expired: .expired
            case .sessionInProgress: .anotherSession   // never fold this into "other"
            default: .other
            }
            let id = ObjectIdentifier(s), onEnd = onEnd
            Task { @MainActor in onEnd?(id, end) }
        }
    }
}