flynn/src/companion/platformClients.ts

import type {
  CompanionEventEnvelope,
  CompanionEventHandler,
  CanvasClearResult,
  CanvasDeleteResult,
  CanvasGetResult,
  CanvasListResult,
  CanvasPutResult,
  CompanionEventPredicate,
  CompanionTypedEventHandler,
  CompanionRuntimeClient,
  DeleteCanvasArtifactInput,
  GetCanvasArtifactInput,
  NodeCapabilitiesResult,
  NodeLocationGetResult,
  NodeLocationSetResult,
  NodeRegisterResult,
  NodeStatusSetResult,
  PutCanvasArtifactInput,
  NodePushTokenSetResult,
  SetNodeLocationInput,
  SystemCapabilitiesResult,
  SystemNodesResult,
} from './runtimeClient.js';
import {
  CompanionHeartbeatLoop,
} from './heartbeatLoop.js';
import type {
  CompanionHeartbeatLoopOptions,
} from './heartbeatLoop.js';

export interface PlatformClientOptions {
  runtime: CompanionRuntimeClient;
  nodeId: string;
  role?: string;
  capabilities?: string[];
  protocolVersion?: number;
  defaultSessionId?: string;
}

export interface RegisterPushTokenInput {
  token: string;
  topic?: string;
  environment?: 'sandbox' | 'production';
}

export type SharedStatusInput = Omit<
  Parameters<CompanionRuntimeClient['setNodeStatus']>[0],
  'platform'
>;

export interface HeartbeatStatusInput {
  appVersion?: SharedStatusInput['appVersion'];
  deviceName?: SharedStatusInput['deviceName'];
  statusText?: SharedStatusInput['statusText'];
  batteryPct?: SharedStatusInput['batteryPct'];
  powerSource?: SharedStatusInput['powerSource'];
}

export interface PlatformBootstrapResult {
  register: NodeRegisterResult;
  capabilities: NodeCapabilitiesResult;
  systemCapabilities?: SystemCapabilitiesResult;
}

export interface PlatformBootstrapOptions {
  includeSystemCapabilities?: boolean;
}

export interface PlatformPutCanvasArtifactInput extends Omit<PutCanvasArtifactInput, 'sessionId'> {
  sessionId?: string;
}

export interface PlatformGetCanvasArtifactInput extends Omit<GetCanvasArtifactInput, 'sessionId'> {
  sessionId?: string;
}

export interface PlatformDeleteCanvasArtifactInput extends Omit<DeleteCanvasArtifactInput, 'sessionId'> {
  sessionId?: string;
}

export class MacOSCompanionClient {
  private readonly runtime: CompanionRuntimeClient;
  private readonly nodeId: string;
  private readonly role: string;
  private readonly capabilities: string[];
  private readonly protocolVersion?: number;
  private readonly defaultSessionId?: string;

  constructor(options: PlatformClientOptions) {
    this.runtime = options.runtime;
    this.nodeId = options.nodeId;
    this.role = options.role ?? 'companion';
    this.capabilities = options.capabilities ?? ['ui.canvas', 'node.location.write', 'node.push.register'];
    this.protocolVersion = options.protocolVersion;
    this.defaultSessionId = options.defaultSessionId;
  }

  connect(): Promise<void> {
    return this.runtime.connect();
  }

  disconnect(): void {
    this.runtime.disconnect();
  }

  dispose(): void {
    this.runtime.dispose();
  }

  register(): Promise<NodeRegisterResult> {
    return this.runtime.registerNode({
      nodeId: this.nodeId,
      role: this.role,
      protocolVersion: this.protocolVersion,
      capabilities: this.capabilities,
    });
  }

  async bootstrap(options?: PlatformBootstrapOptions): Promise<PlatformBootstrapResult> {
    const result = await this.runtime.bootstrapNode(
      {
        nodeId: this.nodeId,
        role: this.role,
        protocolVersion: this.protocolVersion,
        capabilities: this.capabilities,
      },
      { includeSystemCapabilities: options?.includeSystemCapabilities },
    );
    return {
      register: result.register,
      capabilities: result.capabilities,
      systemCapabilities: result.systemCapabilities,
    };
  }

  getCapabilities(): Promise<NodeCapabilitiesResult> {
    return this.runtime.getNodeCapabilities();
  }

  setStatus(status: SharedStatusInput): Promise<NodeStatusSetResult> {
    return this.runtime.setNodeStatus({
      platform: 'macos',
      appVersion: status.appVersion,
      deviceName: status.deviceName,
      statusText: status.statusText,
      batteryPct: status.batteryPct,
      powerSource: status.powerSource,
    });
  }

  publishHeartbeat(input: HeartbeatStatusInput = {}): Promise<NodeStatusSetResult> {
    return this.setStatus({
      appVersion: input.appVersion,
      deviceName: input.deviceName,
      statusText: input.statusText ?? 'heartbeat',
      batteryPct: input.batteryPct,
      powerSource: input.powerSource ?? 'unknown',
    });
  }

  createHeartbeatLoop(options: CompanionHeartbeatLoopOptions = {}): CompanionHeartbeatLoop {
    return new CompanionHeartbeatLoop(this, options);
  }

  setLocation(location: SetNodeLocationInput): Promise<NodeLocationSetResult> {
    return this.runtime.setNodeLocation(location);
  }

  getLocation(): Promise<NodeLocationGetResult> {
    return this.runtime.getNodeLocation();
  }

  registerPushToken(input: RegisterPushTokenInput): Promise<NodePushTokenSetResult> {
    return this.runtime.setNodePushToken({
      provider: 'apns',
      token: input.token,
      topic: input.topic,
      environment: input.environment,
    });
  }

  getSystemCapabilities(): Promise<SystemCapabilitiesResult> {
    return this.runtime.getSystemCapabilities();
  }

  listNodes(): Promise<SystemNodesResult> {
    return this.runtime.listSystemNodes({ platform: 'macos', role: this.role });
  }

  putCanvasArtifact(input: PlatformPutCanvasArtifactInput): Promise<CanvasPutResult> {
    return this.runtime.putCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
      type: input.type,
      title: input.title,
      content: input.content,
      metadata: input.metadata,
    });
  }

  getCanvasArtifact(input: PlatformGetCanvasArtifactInput): Promise<CanvasGetResult> {
    return this.runtime.getCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
    });
  }

  listCanvasArtifacts(sessionId?: string): Promise<CanvasListResult> {
    return this.runtime.listCanvasArtifacts(this.resolveSessionId(sessionId));
  }

  deleteCanvasArtifact(input: PlatformDeleteCanvasArtifactInput): Promise<CanvasDeleteResult> {
    return this.runtime.deleteCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
    });
  }

  clearCanvasArtifacts(sessionId?: string): Promise<CanvasClearResult> {
    return this.runtime.clearCanvasArtifacts(this.resolveSessionId(sessionId));
  }

  subscribeAgentStream<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeAgentStream(handler);
  }

  subscribeAgentTyping<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeAgentTyping(handler);
  }

  subscribeContextWarning<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeContextWarning(handler);
  }

  waitForAgentStream<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForAgentStream(options);
  }

  subscribeEvents(handler: CompanionEventHandler): () => void {
    return this.runtime.subscribeEvents(handler);
  }

  waitForAnyEvent<TData = unknown>(
    eventNames: readonly string[],
    options?: {
      timeoutMs?: number;
      predicate?: (event: string, data: TData) => boolean;
      signal?: AbortSignal;
    },
  ): Promise<CompanionEventEnvelope<TData>> {
    return this.runtime.waitForAnyEvent(eventNames, options);
  }

  waitForAgentTyping<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForAgentTyping(options);
  }

  waitForContextWarning<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForContextWarning(options);
  }

  private resolveSessionId(sessionId?: string): string {
    const resolved = sessionId ?? this.defaultSessionId;
    if (!resolved) {
      throw new Error('sessionId is required (provide one or configure defaultSessionId)');
    }
    return resolved;
  }
}

export class IOSCompanionClient {
  private readonly runtime: CompanionRuntimeClient;
  private readonly nodeId: string;
  private readonly role: string;
  private readonly capabilities: string[];
  private readonly protocolVersion?: number;
  private readonly defaultSessionId?: string;

  constructor(options: PlatformClientOptions) {
    this.runtime = options.runtime;
    this.nodeId = options.nodeId;
    this.role = options.role ?? 'companion';
    this.capabilities = options.capabilities ?? ['node.location.write', 'node.push.register'];
    this.protocolVersion = options.protocolVersion;
    this.defaultSessionId = options.defaultSessionId;
  }

  connect(): Promise<void> {
    return this.runtime.connect();
  }

  disconnect(): void {
    this.runtime.disconnect();
  }

  dispose(): void {
    this.runtime.dispose();
  }

  register(): Promise<NodeRegisterResult> {
    return this.runtime.registerNode({
      nodeId: this.nodeId,
      role: this.role,
      protocolVersion: this.protocolVersion,
      capabilities: this.capabilities,
    });
  }

  async bootstrap(options?: PlatformBootstrapOptions): Promise<PlatformBootstrapResult> {
    const result = await this.runtime.bootstrapNode(
      {
        nodeId: this.nodeId,
        role: this.role,
        protocolVersion: this.protocolVersion,
        capabilities: this.capabilities,
      },
      { includeSystemCapabilities: options?.includeSystemCapabilities },
    );
    return {
      register: result.register,
      capabilities: result.capabilities,
      systemCapabilities: result.systemCapabilities,
    };
  }

  getCapabilities(): Promise<NodeCapabilitiesResult> {
    return this.runtime.getNodeCapabilities();
  }

  setStatus(status: SharedStatusInput): Promise<NodeStatusSetResult> {
    return this.runtime.setNodeStatus({
      platform: 'ios',
      appVersion: status.appVersion,
      deviceName: status.deviceName,
      statusText: status.statusText,
      batteryPct: status.batteryPct,
      powerSource: status.powerSource,
    });
  }

  publishHeartbeat(input: HeartbeatStatusInput = {}): Promise<NodeStatusSetResult> {
    return this.setStatus({
      appVersion: input.appVersion,
      deviceName: input.deviceName,
      statusText: input.statusText ?? 'heartbeat',
      batteryPct: input.batteryPct,
      powerSource: input.powerSource ?? 'unknown',
    });
  }

  createHeartbeatLoop(options: CompanionHeartbeatLoopOptions = {}): CompanionHeartbeatLoop {
    return new CompanionHeartbeatLoop(this, options);
  }

  setLocation(location: SetNodeLocationInput): Promise<NodeLocationSetResult> {
    return this.runtime.setNodeLocation(location);
  }

  getLocation(): Promise<NodeLocationGetResult> {
    return this.runtime.getNodeLocation();
  }

  registerPushToken(input: RegisterPushTokenInput): Promise<NodePushTokenSetResult> {
    return this.runtime.setNodePushToken({
      provider: 'apns',
      token: input.token,
      topic: input.topic,
      environment: input.environment,
    });
  }

  getSystemCapabilities(): Promise<SystemCapabilitiesResult> {
    return this.runtime.getSystemCapabilities();
  }

  listNodes(): Promise<SystemNodesResult> {
    return this.runtime.listSystemNodes({ platform: 'ios', role: this.role });
  }

  putCanvasArtifact(input: PlatformPutCanvasArtifactInput): Promise<CanvasPutResult> {
    return this.runtime.putCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
      type: input.type,
      title: input.title,
      content: input.content,
      metadata: input.metadata,
    });
  }

  getCanvasArtifact(input: PlatformGetCanvasArtifactInput): Promise<CanvasGetResult> {
    return this.runtime.getCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
    });
  }

  listCanvasArtifacts(sessionId?: string): Promise<CanvasListResult> {
    return this.runtime.listCanvasArtifacts(this.resolveSessionId(sessionId));
  }

  deleteCanvasArtifact(input: PlatformDeleteCanvasArtifactInput): Promise<CanvasDeleteResult> {
    return this.runtime.deleteCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
    });
  }

  clearCanvasArtifacts(sessionId?: string): Promise<CanvasClearResult> {
    return this.runtime.clearCanvasArtifacts(this.resolveSessionId(sessionId));
  }

  subscribeAgentStream<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeAgentStream(handler);
  }

  subscribeAgentTyping<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeAgentTyping(handler);
  }

  subscribeContextWarning<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeContextWarning(handler);
  }

  waitForAgentStream<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForAgentStream(options);
  }

  subscribeEvents(handler: CompanionEventHandler): () => void {
    return this.runtime.subscribeEvents(handler);
  }

  waitForAnyEvent<TData = unknown>(
    eventNames: readonly string[],
    options?: {
      timeoutMs?: number;
      predicate?: (event: string, data: TData) => boolean;
      signal?: AbortSignal;
    },
  ): Promise<CompanionEventEnvelope<TData>> {
    return this.runtime.waitForAnyEvent(eventNames, options);
  }

  waitForAgentTyping<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForAgentTyping(options);
  }

  waitForContextWarning<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForContextWarning(options);
  }

  private resolveSessionId(sessionId?: string): string {
    const resolved = sessionId ?? this.defaultSessionId;
    if (!resolved) {
      throw new Error('sessionId is required (provide one or configure defaultSessionId)');
    }
    return resolved;
  }
}

export class AndroidCompanionClient {
  private readonly runtime: CompanionRuntimeClient;
  private readonly nodeId: string;
  private readonly role: string;
  private readonly capabilities: string[];
  private readonly protocolVersion?: number;
  private readonly defaultSessionId?: string;

  constructor(options: PlatformClientOptions) {
    this.runtime = options.runtime;
    this.nodeId = options.nodeId;
    this.role = options.role ?? 'companion';
    this.capabilities = options.capabilities ?? ['node.location.write', 'node.push.register'];
    this.protocolVersion = options.protocolVersion;
    this.defaultSessionId = options.defaultSessionId;
  }

  connect(): Promise<void> {
    return this.runtime.connect();
  }

  disconnect(): void {
    this.runtime.disconnect();
  }

  dispose(): void {
    this.runtime.dispose();
  }

  register(): Promise<NodeRegisterResult> {
    return this.runtime.registerNode({
      nodeId: this.nodeId,
      role: this.role,
      protocolVersion: this.protocolVersion,
      capabilities: this.capabilities,
    });
  }

  async bootstrap(options?: PlatformBootstrapOptions): Promise<PlatformBootstrapResult> {
    const result = await this.runtime.bootstrapNode(
      {
        nodeId: this.nodeId,
        role: this.role,
        protocolVersion: this.protocolVersion,
        capabilities: this.capabilities,
      },
      { includeSystemCapabilities: options?.includeSystemCapabilities },
    );
    return {
      register: result.register,
      capabilities: result.capabilities,
      systemCapabilities: result.systemCapabilities,
    };
  }

  getCapabilities(): Promise<NodeCapabilitiesResult> {
    return this.runtime.getNodeCapabilities();
  }

  setStatus(status: SharedStatusInput): Promise<NodeStatusSetResult> {
    return this.runtime.setNodeStatus({
      platform: 'android',
      appVersion: status.appVersion,
      deviceName: status.deviceName,
      statusText: status.statusText,
      batteryPct: status.batteryPct,
      powerSource: status.powerSource,
    });
  }

  publishHeartbeat(input: HeartbeatStatusInput = {}): Promise<NodeStatusSetResult> {
    return this.setStatus({
      appVersion: input.appVersion,
      deviceName: input.deviceName,
      statusText: input.statusText ?? 'heartbeat',
      batteryPct: input.batteryPct,
      powerSource: input.powerSource ?? 'unknown',
    });
  }

  createHeartbeatLoop(options: CompanionHeartbeatLoopOptions = {}): CompanionHeartbeatLoop {
    return new CompanionHeartbeatLoop(this, options);
  }

  setLocation(location: SetNodeLocationInput): Promise<NodeLocationSetResult> {
    return this.runtime.setNodeLocation(location);
  }

  getLocation(): Promise<NodeLocationGetResult> {
    return this.runtime.getNodeLocation();
  }

  registerPushToken(token: string): Promise<NodePushTokenSetResult> {
    return this.runtime.setNodePushToken({
      provider: 'fcm',
      token,
    });
  }

  getSystemCapabilities(): Promise<SystemCapabilitiesResult> {
    return this.runtime.getSystemCapabilities();
  }

  listNodes(): Promise<SystemNodesResult> {
    return this.runtime.listSystemNodes({ platform: 'android', role: this.role });
  }

  putCanvasArtifact(input: PlatformPutCanvasArtifactInput): Promise<CanvasPutResult> {
    return this.runtime.putCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
      type: input.type,
      title: input.title,
      content: input.content,
      metadata: input.metadata,
    });
  }

  getCanvasArtifact(input: PlatformGetCanvasArtifactInput): Promise<CanvasGetResult> {
    return this.runtime.getCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
    });
  }

  listCanvasArtifacts(sessionId?: string): Promise<CanvasListResult> {
    return this.runtime.listCanvasArtifacts(this.resolveSessionId(sessionId));
  }

  deleteCanvasArtifact(input: PlatformDeleteCanvasArtifactInput): Promise<CanvasDeleteResult> {
    return this.runtime.deleteCanvasArtifact({
      sessionId: this.resolveSessionId(input.sessionId),
      artifactId: input.artifactId,
    });
  }

  clearCanvasArtifacts(sessionId?: string): Promise<CanvasClearResult> {
    return this.runtime.clearCanvasArtifacts(this.resolveSessionId(sessionId));
  }

  subscribeAgentStream<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeAgentStream(handler);
  }

  subscribeAgentTyping<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeAgentTyping(handler);
  }

  subscribeContextWarning<TData = unknown>(
    handler: CompanionTypedEventHandler<TData>,
  ): () => void {
    return this.runtime.subscribeContextWarning(handler);
  }

  waitForAgentStream<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForAgentStream(options);
  }

  subscribeEvents(handler: CompanionEventHandler): () => void {
    return this.runtime.subscribeEvents(handler);
  }

  waitForAnyEvent<TData = unknown>(
    eventNames: readonly string[],
    options?: {
      timeoutMs?: number;
      predicate?: (event: string, data: TData) => boolean;
      signal?: AbortSignal;
    },
  ): Promise<CompanionEventEnvelope<TData>> {
    return this.runtime.waitForAnyEvent(eventNames, options);
  }

  waitForAgentTyping<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForAgentTyping(options);
  }

  waitForContextWarning<TData = unknown>(options?: {
    timeoutMs?: number;
    predicate?: CompanionEventPredicate<TData>;
    signal?: AbortSignal;
  }): Promise<TData> {
    return this.runtime.waitForContextWarning(options);
  }

  private resolveSessionId(sessionId?: string): string {
    const resolved = sessionId ?? this.defaultSessionId;
    if (!resolved) {
      throw new Error('sessionId is required (provide one or configure defaultSessionId)');
    }
    return resolved;
  }
}