flynn/apps/companion/ios/IOSCompanionRuntime.swift

import Foundation

enum CompanionRuntimeError: Error {
  case invalidGatewayURL
  case notConnected
  case encodingFailure
}

private struct JsonRpcRequest<Params: Encodable>: Encodable {
  let jsonrpc = "2.0"
  let id: Int
  let method: String
  let params: Params
}

struct CompanionStatusPayload: Encodable {
  let platform: String
  let appVersion: String?
  let statusText: String?
  let batteryPct: Double?
}

struct CompanionLocationPayload: Encodable {
  let latitude: Double
  let longitude: Double
  let source: String
}

private struct NodeRegisterParams: Encodable {
  let nodeId: String
  let role: String
  let platform: String
  let capabilities: [String]
}

private struct NodeStatusSetParams: Encodable {
  let nodeId: String
  let status: CompanionStatusPayload
}

private struct NodeLocationSetParams: Encodable {
  let nodeId: String
  let location: CompanionLocationPayload
}

private struct NodePushTokenSetParams: Encodable {
  let nodeId: String
  let provider: String
  let token: String
  let topic: String?
  let environment: String?
}

private struct AgentSendParams: Encodable {
  let message: String
  let sessionId: String?
  let awaitResponse: Bool
}

final class IOSCompanionRuntime {
  private let bootstrap: CompanionBootstrap
  private let session: URLSession
  private var socket: URLSessionWebSocketTask?
  private var nextRequestId = 1
  private var heartbeatLoop: Task<Void, Never>?

  init(bootstrap: CompanionBootstrap, session: URLSession = .shared) {
    self.bootstrap = bootstrap
    self.session = session
  }

  func connect() throws {
    guard socket == nil else { return }
    guard let url = URL(string: bootstrap.gateway.url) else {
      throw CompanionRuntimeError.invalidGatewayURL
    }
    var request = URLRequest(url: url)
    if let token = bootstrap.gateway.token, !token.isEmpty {
      request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
    }
    let task = session.webSocketTask(with: request)
    task.resume()
    socket = task
  }

  func disconnect() {
    heartbeatLoop?.cancel()
    heartbeatLoop = nil
    socket?.cancel(with: .goingAway, reason: nil)
    socket = nil
  }

  func registerNode() async throws {
    try await sendRpc(
      method: "node.register",
      params: NodeRegisterParams(
        nodeId: bootstrap.node.nodeId,
        role: bootstrap.node.role,
        platform: bootstrap.node.platform,
        capabilities: bootstrap.node.capabilities
      )
    )
  }

  func publishStatus(_ status: CompanionStatusPayload) async throws {
    try await sendRpc(
      method: "node.status.set",
      params: NodeStatusSetParams(nodeId: bootstrap.node.nodeId, status: status)
    )
  }

  func publishLocation(_ location: CompanionLocationPayload) async throws {
    try await sendRpc(
      method: "node.location.set",
      params: NodeLocationSetParams(nodeId: bootstrap.node.nodeId, location: location)
    )
  }

  func registerPushToken(token: String, topic: String? = nil, environment: String? = nil) async throws {
    try await sendRpc(
      method: "node.push_token.set",
      params: NodePushTokenSetParams(
        nodeId: bootstrap.node.nodeId,
        provider: "apns",
        token: token,
        topic: topic,
        environment: environment
      )
    )
  }

  func sendHandoff(message: String, sessionId: String? = nil) async throws {
    try await sendRpc(
      method: "agent.send",
      params: AgentSendParams(message: message, sessionId: sessionId, awaitResponse: true)
    )
  }

  func startHeartbeatLoop(statusFactory: @escaping () -> CompanionStatusPayload) {
    heartbeatLoop?.cancel()
    let interval = UInt64(max(bootstrap.runtime.heartbeatSeconds, 1)) * 1_000_000_000
    heartbeatLoop = Task { [weak self] in
      while !Task.isCancelled {
        guard let self else { return }
        try? await self.publishStatus(statusFactory())
        try? await Task.sleep(nanoseconds: interval)
      }
    }
  }

  private func sendRpc<Params: Encodable>(method: String, params: Params) async throws {
    guard let socket else {
      throw CompanionRuntimeError.notConnected
    }
    let request = JsonRpcRequest(id: nextRequestId, method: method, params: params)
    nextRequestId += 1
    let data = try JSONEncoder().encode(request)
    guard let text = String(data: data, encoding: .utf8) else {
      throw CompanionRuntimeError.encodingFailure
    }
    try await socket.send(.string(text))
  }
}

// Generated for node: ios-reference-shell (ios)