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
    private var lifecycleTask: Task<Void, Never>?
    private var pending: [UUID: RelayReplyOnce] = [:]
    private var latestContext: RelayApplicationContext?

    public func publishContext(_ context: RelayApplicationContext) {
        latestContext = context
        flushContext()
    }

    private func flushContext() {
        guard session.activationState == .activated, let latestContext else { return }
        if let dictionary = try? latestContext.dictionary() {
            try? session.updateApplicationContext(dictionary)
        }
    }

    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, session.activationState == .activated else {
            didDeactivate()
            state = .unavailable
            activationReady = false
            return
        }
        guard peer == nil else { flushContext(); return }
        let newPeer = WatchPeerGeneration()
        peer = newPeer
        activationReady = false
        state = .ready(reachable: session.isReachable)
        let prior = lifecycleTask
        lifecycleTask = Task { [weak self] in
            await prior?.value
            await self?.coordinator.watchBridgeDidActivate(peer: newPeer)
            guard let self, self.peer == newPeer else { return }
            self.activationReady = true
            self.flushContext()
        }
    }

    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
        for reply in pending.values { reply.reply(relayErrorReply(RelayFailure(code: .unavailable))) }
        let prior = lifecycleTask
        lifecycleTask = Task {
            await prior?.value
            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 inbound = RelayInboundMessage(message: message, replyHandler: replyHandler)
        guard pending.count < 16 else {
            inbound.reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
            return
        }
        let receivedPeer = peer
        let id = UUID()
        pending[id] = inbound.reply
        Task { @MainActor [weak self] in
            guard let self else {
                inbound.reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
                return
            }
            guard self.peer == receivedPeer, receivedPeer != nil else {
                inbound.reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
                self.pending.removeValue(forKey: id)
                return
            }
            await self.route(message: inbound.message, reply: inbound.reply)
            self.pending.removeValue(forKey: id)
        }
        Task { @MainActor [weak self] in
            try? await Task.sleep(for: .seconds(1))
            guard let self, let reply = self.pending[id] else { return }
            reply.reply(relayErrorReply(RelayFailure(code: .unavailable)))
        }
    }

    private func route(message: [String: Any], reply: RelayReplyOnce) async {
        guard activationReady, session.activationState == .activated, 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
            }
            let result = await ({
                await self.coordinator.beginWatchSession(request, peer: currentPeer)
            })()
            guard peer == currentPeer, activationReady 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
            }
            let result = await ({
                await self.coordinator.endWatchSession(request, peer: currentPeer)
            })()
            guard peer == currentPeer, activationReady 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
            }
            let receipt = await ({
                await self.coordinator.submitWatchCommand(envelope, peer: currentPeer)
            })()
            guard peer == currentPeer, activationReady 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))
    }

}