wristivo/Adapters/WatchBridge/Sources/Watch/WatchRelayBridge.swift

import Foundation
import RemoteCore
import WatchBridgeCommon

@MainActor
public final class WatchRelayBridge: NSObject {
    public enum State: Sendable, Equatable {
        case inactive
        case activating
        case unavailable
        case ready(reachable: Bool)
        case sessionNegotiating
        case armed(sessionID: UUID, televisionID: UUID)
    }

    public private(set) var state: State = .inactive {
        didSet { onStateChange?(state) }
    }
    public var onStateChange: ((State) -> Void)?

    private let session: any RelaySessionSending
    public private(set) var grant: RemoteSessionGrant?
    private var generation = UUID()
    private var nextSequence: UInt64 = 0
    private var activationReady = false
    private var pendingReplies: [UUID: RelayReplyOnce] = [:]

    public init(session: any RelaySessionSending) {
        self.session = session
        super.init()
    }

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

    public func activationDidComplete(error: Error?) {
        guard error == nil else {
            activationReady = false
            state = .unavailable
            return
        }
        activationReady = true
        state = .ready(reachable: session.isReachable)
    }

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

    public func didBecomeInactive() {
        activationReady = false
        invalidatePending()
        state = .inactive
    }

    public func didDeactivate() {
        activationReady = false
        grant = nil
        nextSequence = 0
        invalidatePending()
        state = .inactive
    }

    public func beginSession() {
        guard grant == nil, state != .sessionNegotiating else { return }
        guard activationReady, session.isReachable, grant == nil else {
            state = .unavailable
            return
        }
        invalidatePending()
        let requestGeneration = generation
        state = .sessionNegotiating
        let requestID = UUID()
        let request = BeginRemoteSessionRequest()
        guard let message = try? RelayContainer(kind: .beginSession, value: request).dictionary() else {
            state = .unavailable
            return
        }
        send(message: message, requestID: requestID) { [weak self] response in
            guard let self, self.generation == requestGeneration, self.session.isReachable else { return }
            guard let container = try? RelayContainer(dictionary: response),
                  container.kind == .session,
                  let received = try? container.decode(RemoteSessionGrant.self,
                                                        expectedKind: .session) else {
                self.state = .unavailable
                return
            }
            self.grant = received
            self.nextSequence = 0
            self.state = .armed(sessionID: received.sessionID,
                                televisionID: received.televisionID)
        }
    }

    public func endSession() {
        let oldGrant = grant
        invalidatePending()
        state = activationReady ? .ready(reachable: session.isReachable) : .inactive
        guard activationReady, session.isReachable, let oldGrant,
              let message = try? RelayContainer(kind: .endSession, value: EndRemoteSessionRequest(sessionID: oldGrant.sessionID)).dictionary() else { return }
        // Locally revoke immediately; a delayed end reply cannot revoke a newer grant.
        send(message: message, requestID: UUID()) { _ in }
    }

    public func sendCommandResult(_ command: RemoteCommand) async -> Result<CommandReceipt, RelayFailure> {
        await withCheckedContinuation { continuation in
            sendCommand(command) { continuation.resume(returning: $0) }
        }
    }

    /// One command is one live sendMessage. A sequence is consumed before the
    /// send; ambiguous timeout/failure is never replayed.
    public func sendCommand(_ command: RemoteCommand, now: Date = Date(),
                            completion: @escaping @MainActor (Result<CommandReceipt, RelayFailure>) -> Void = { _ in }) {
        guard activationReady, session.isReachable, let grant else {
            state = .unavailable
            completion(.failure(.init(code: .unavailable)))
            return
        }
        guard nextSequence < UInt64.max else {
            self.grant = nil
            state = .unavailable
            completion(.failure(.init(code: .sessionInvalidated)))
            return
        }
        let sequence = nextSequence
        nextSequence += 1
        let lifetime = Double(grant.maximumCommandLifetimeMilliseconds) / 1_000
        let envelope = CommandEnvelope(sessionID: grant.sessionID,
                                       televisionID: grant.televisionID,
                                       sequence: sequence,
                                       issuedAt: now,
                                       expiresAt: now.addingTimeInterval(lifetime),
                                       command: command)
        guard let message = try? RelayContainer(kind: .command, value: envelope).dictionary() else {
            completion(.failure(.init(code: .malformedRequest)))
            return
        }
        let requestGeneration = generation
        send(message: message, requestID: UUID()) { [weak self] response in
            guard let self, self.generation == requestGeneration,
                  let container = try? RelayContainer(dictionary: response),
                  container.kind == .receipt,
                  let receipt = try? container.decode(CommandReceipt.self,
                                                       expectedKind: .receipt),
                  receipt.id == envelope.id else { completion(.failure(.init(code: .unavailable))); return }
            completion(.success(receipt))
            // A receipt is command feedback, not evidence that the TV changed.
            if receipt.status != .sentToTransport { self.invalidatePending(); self.state = .unavailable }
        }
    }

    private func send(message: [String: Any], requestID: UUID,
                      response: @escaping @MainActor ([String: Any]) -> Void) {
        guard pendingReplies.count < 32 else {
            state = .unavailable
            response(relayErrorReply(RelayFailure(code: .unavailable)))
            return
        }
        let once = RelayReplyOnce { [weak self] value in
            Task { @MainActor [weak self] in
                guard let self, self.pendingReplies.removeValue(forKey: requestID) != nil else { return }
                response(value)
            }
        }
        pendingReplies[requestID] = once
        session.sendMessage(message, replyHandler: { value in
            once.reply(value)
        }, errorHandler: { _ in
            once.reply(relayErrorReply(RelayFailure(code: .unavailable)))
        })
        Task { @MainActor [weak self] in
            try? await Task.sleep(for: .milliseconds(1_000))
            self?.pendingReplies[requestID]?.reply(relayErrorReply(RelayFailure(code: .unavailable)))
        }
    }

    private func invalidatePending() {
        generation = UUID()
        let replies = Array(pendingReplies.values)
        for reply in replies { reply.reply(relayErrorReply(RelayFailure(code: .sessionInvalidated))) }
        grant = nil
        nextSequence = 0
    }
}