wristivo/Adapters/WatchBridge/Sources/iPhone/PhoneRelayBridge.swift

import Foundation
import RemoteCore
import WatchBridgeCommon

@MainActor
public final class PhoneRelayBridge: NSObject {
    public enum State: Sendable, Equatable {
        case inactive
        case activating
        case ready(reachable: Bool)
        case unavailable
    }

    public private(set) var state: State = .inactive
    private let session: any RelaySessionSending
    private let coordinator: any PhoneCommandCoordinating
    private var peer: WatchPeerGeneration?
    private var activationReady = false

    public init(session: any RelaySessionSending,
                dependency: PhoneRelayDependency) {
        self.session = session
        self.coordinator = dependency.coordinator
    }

    public func activate() {
        guard peer == nil else { return }
        state = .activating
        activationReady = false
        session.activate()
    }

    public func activationDidComplete(error: Error?) {
        guard error == nil else {
            state = .unavailable
            activationReady = false
            return
        }
        let newPeer = WatchPeerGeneration()
        peer = newPeer
        activationReady = false
        state = .ready(reachable: session.isReachable)
        Task { [weak self] in
            await self?.coordinator.watchBridgeDidActivate(peer: newPeer)
            guard let self, self.peer == newPeer else { return }
            self.activationReady = true
        }
    }

    public func reachabilityDidChange() {
        guard peer != nil, activationReady else { return }
        state = .ready(reachable: session.isReachable)
    }

    public func didBecomeInactive() {
        // iOS will normally follow with sessionDidDeactivate, but invalidate
        // immediately so no command can slip through during counterpart swap.
        didDeactivate()
    }

    public func didDeactivate() {
        guard let oldPeer = peer else { return }
        peer = nil
        activationReady = false
        state = .inactive
        Task { await coordinator.watchBridgeDidDeactivate(peer: oldPeer) }
    }

    /// Entry point for the native WCSession delegate. Every path completes the
    /// supplied reply exactly once, including malformed, unavailable and timed
    /// out requests.
    public func receive(message: [String: Any],
                        replyHandler: @escaping ([String: Any]) -> Void) {
        let reply = RelayReplyOnce(replyHandler)
        Task { @MainActor [weak self] in
            guard let self else {
                reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
                return
            }
            await self.route(message: message, reply: reply)
        }
    }

    private func route(message: [String: Any], reply: RelayReplyOnce) async {
        guard activationReady, let currentPeer = peer else {
            reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
            return
        }
        let container: RelayContainer
        do {
            container = try RelayContainer(dictionary: message)
        } catch {
            reply.reply(relayErrorReply(RelayFailure(code: .malformedRequest)))
            return
        }

        switch container.kind {
        case .beginSession:
            guard let request = try? container.decode(BeginRemoteSessionRequest.self,
                                                       expectedKind: .beginSession) else {
                reply.reply(relayErrorReply(RelayFailure(code: .malformedRequest)))
                return
            }
            guard let result = await withTimeout({
                await self.coordinator.beginWatchSession(request, peer: currentPeer)
            }) else {
                reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
                return
            }
            switch result {
            case .success(let grant):
                reply.reply(Self.reply(kind: .session, value: grant))
            case .failure(let failure):
                reply.reply(relayErrorReply(failure))
            }

        case .endSession:
            guard let request = try? container.decode(EndRemoteSessionRequest.self,
                                                       expectedKind: .endSession) else {
                reply.reply(relayErrorReply(RelayFailure(code: .malformedRequest)))
                return
            }
            guard let result = await withTimeout({
                await self.coordinator.endWatchSession(request, peer: currentPeer)
            }) else {
                reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
                return
            }
            switch result {
            case .success(let ended):
                reply.reply(Self.reply(kind: .ended, value: ended))
            case .failure(let failure):
                reply.reply(relayErrorReply(failure))
            }

        case .command:
            guard let envelope = try? container.decode(CommandEnvelope.self,
                                                        expectedKind: .command) else {
                reply.reply(relayErrorReply(RelayFailure(code: .malformedRequest)))
                return
            }
            guard let receipt = await withTimeout({
                await self.coordinator.submitWatchCommand(envelope, peer: currentPeer)
            }) else {
                reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
                return
            }
            guard receipt.id == envelope.id else {
                reply.reply(relayErrorReply(RelayFailure(code: .malformedRequest)))
                return
            }
            reply.reply(Self.reply(kind: .receipt, value: receipt))

        case .session, .ended, .receipt, .error:
            reply.reply(relayErrorReply(RelayFailure(code: .malformedRequest)))
        }
    }

    private static func reply<Value: Encodable>(kind: RelayMessageKind,
                                                 value: Value) -> [String: Any] {
        (try? RelayContainer(kind: kind, value: value).dictionary()) ??
            relayErrorReply(RelayFailure(code: .malformedRequest))
    }

    private func withTimeout<Value: Sendable>(
        _ operation: @escaping @Sendable () async -> Value
    ) async -> Value? {
        await withTaskGroup(of: Value?.self) { group in
            group.addTask { await operation() }
            group.addTask {
                try? await Task.sleep(for: .milliseconds(1_000))
                return nil
            }
            let result = await group.next() ?? nil
            group.cancelAll()
            return result
        }
    }
}