rainbow-dragon/Sources/ForgeCamera/ConnectionSupervisor.swift

import Foundation

/// Exponential backoff with cap.
public struct BackoffPolicy: Sendable {
    public let initial: Duration
    public let multiplier: Double
    public let maxDelay: Duration
    private var current: Duration?

    public init(initial: Duration, multiplier: Double, maxDelay: Duration) {
        self.initial = initial
        self.multiplier = multiplier
        self.maxDelay = maxDelay
    }

    public mutating func nextDelay() -> Duration {
        guard let c = current else {
            current = initial
            return initial
        }
        let scaled = Duration.seconds(Double(c.components.seconds) * multiplier
            + Double(c.components.attoseconds) * multiplier / 1e18)
        let next = scaled > maxDelay ? maxDelay : scaled
        current = next
        return next
    }

    public mutating func reset() {
        current = nil
    }
}

/// Owns connect/retry/health for one driver. Republishes driver states.
public actor ConnectionSupervisor {
    private let driver: any CameraDriver
    private var backoff: BackoffPolicy
    private var connectTask: Task<Void, Never>?
    private var pumpTask: Task<Void, Never>?

    public private(set) var health: ConnectionStatus = .disconnected

    /// Subscriber continuations behind a lock — registration is synchronous
    /// at subscription time, so states emitted immediately after `states`
    /// is accessed are never dropped (async actor-hop registration raced).
    private final class SubscriberBox: @unchecked Sendable {
        let lock = NSLock()
        var continuations: [UUID: AsyncStream<CameraState>.Continuation] = [:]

        func add(_ id: UUID, _ c: AsyncStream<CameraState>.Continuation) {
            lock.lock()
            continuations[id] = c
            lock.unlock()
        }

        func remove(_ id: UUID) {
            lock.lock()
            continuations.removeValue(forKey: id)
            lock.unlock()
        }

        func yieldAll(_ s: CameraState) {
            lock.lock()
            let conts = Array(continuations.values)
            lock.unlock()
            for c in conts { c.yield(s) }
        }
    }

    private nonisolated let subscribers = SubscriberBox()
    private var connectionWaiters: [UUID] = []
    private var waiterMap: [UUID: CheckedContinuation<Void, Never>] = [:]

    public init(driver: any CameraDriver, backoff: BackoffPolicy) {
        self.driver = driver
        self.backoff = backoff
    }

    /// New subscription to republished states.
    public nonisolated var states: AsyncStream<CameraState> {
        AsyncStream { continuation in
            let id = UUID()
            subscribers.add(id, continuation)
            continuation.onTermination = { [subscribers] _ in
                subscribers.remove(id)
            }
        }
    }

    public func start() throws {
        guard connectTask == nil else { return }
        health = .connecting

        // Pump driver states into republish + health.
        pumpTask = Task { [driver] in
            for await s in driver.state {
                await self.handle(state: s)
            }
        }

        // Connect with retry.
        connectTask = Task {
            while !Task.isCancelled {
                do {
                    try await driver.connect()
                    return
                } catch {
                    let delay = self.nextBackoffDelay()
                    try? await Task.sleep(for: delay)
                }
            }
        }
    }

    private func nextBackoffDelay() -> Duration {
        backoff.nextDelay()
    }

    private func handle(state: CameraState) {
        health = state.connection
        if state.connection == .connected {
            backoff.reset()
            resumeAllWaiters()
        }
        subscribers.yieldAll(state)
    }

    /// Await connected health. Returns false on timeout.
    /// Continuation registered synchronously on the actor — no gap for a
    /// connected event to slip through unobserved. Timeout resumes the
    /// continuation directly (cancellation can't interrupt a checked continuation).
    public func waitForConnection(timeout: Duration) async -> Bool {
        if health == .connected { return true }
        await withCheckedContinuation { (c: CheckedContinuation<Void, Never>) in
            let id = UUID()
            waiterMap[id] = c
            connectionWaiters.append(id)
            Task {
                try? await Task.sleep(for: timeout)
                self.timeoutWaiter(id: id)
            }
        }
        return health == .connected
    }

    private func timeoutWaiter(id: UUID) {
        guard let c = waiterMap.removeValue(forKey: id) else { return }
        connectionWaiters.removeAll { $0 == id }
        c.resume()
    }

    public func stop() async {
        connectTask?.cancel()
        pumpTask?.cancel()
        connectTask = nil
        pumpTask = nil
        await driver.disconnect()
        health = .disconnected
        resumeAllWaiters()
    }

    private func resumeAllWaiters() {
        let ids = connectionWaiters
        connectionWaiters.removeAll()
        for id in ids {
            waiterMap.removeValue(forKey: id)?.resume()
        }
    }
}