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

    private var stateContinuations: [UUID: AsyncStream<CameraState>.Continuation] = [:]
    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()
            Task { await self.register(id: id, continuation: continuation) }
            continuation.onTermination = { _ in
                Task { await self.unregister(id: id) }
            }
        }
    }

    private func register(id: UUID, continuation: AsyncStream<CameraState>.Continuation) {
        stateContinuations[id] = continuation
    }

    private func unregister(id: UUID) {
        stateContinuations.removeValue(forKey: 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()
        }
        for c in stateContinuations.values {
            c.yield(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()
        }
    }
}