src/lib/elements/robot/drivers/RemoteConsumer.ts

import type { Consumer, ConnectionStatus, RobotCommand, RemoteDriverConfig } from "../models.js";
import { robotics } from "@robothub/transport-server-client";
import type { JointData } from "@robothub/transport-server-client/robotics";

export class RemoteConsumer implements Consumer {
	readonly id: string;
	readonly name: string;
	readonly config: RemoteDriverConfig;

	private _status: ConnectionStatus = { isConnected: false };
	private statusCallbacks: ((status: ConnectionStatus) => void)[] = [];
	private commandCallbacks: ((command: RobotCommand) => void)[] = [];

	private consumer: robotics.RoboticsConsumer | null = null;
	private client: robotics.RoboticsClientCore | null = null;
	private workspaceId: string | null = null;

	constructor(config: RemoteDriverConfig) {
		this.config = config;
		this.id = `remote-consumer-${config.robotId}-${Date.now()}`;
		this.name = `Remote Consumer (${config.robotId})`;
	}

	get status(): ConnectionStatus {
		return this._status;
	}

	async connect(joinExistingRoom = false): Promise<void> {
		try {
			const serverUrl = this.config.url.replace(/^ws/, "http");
			console.log(
				`[RemoteConsumer] Connecting to ${serverUrl} for robot ${this.config.robotId} (join mode: ${joinExistingRoom})`
			);

			// Create core client for room management
			this.client = new robotics.RoboticsClientCore(serverUrl);

			// Create consumer to receive commands
			this.consumer = new robotics.RoboticsConsumer(serverUrl);

			// Set up event handlers
			this.consumer.onConnected(() => {
				console.log(`[RemoteConsumer] Connected to room ${this.config.robotId}`);
			});

			this.consumer.onDisconnected(() => {
				console.log(`[RemoteConsumer] Disconnected from room ${this.config.robotId}`);
			});

			this.consumer.onError((error: string) => {
				console.error(`[RemoteConsumer] Error:`, error);
				this._status = { isConnected: false, error: `Consumer error: ${error}` };
				this.notifyStatusChange();
			});

			// RECEIVE joint updates and forward as normalized commands
			this.consumer.onJointUpdate((joints: JointData[]) => {
				console.debug(`[RemoteConsumer] Received joint update:`, joints);

				const command: RobotCommand = {
					timestamp: Date.now(),
					joints: joints.map((joint: JointData) => ({
						name: joint.name,
						value: joint.value // Already normalized from server
					}))
				};
				this.notifyCommand(command);
			});

			// RECEIVE state sync
			this.consumer.onStateSync((state: Record<string, number>) => {
				console.debug(`[RemoteConsumer] Received state sync:`, state);

				const joints = Object.entries(state).map(([name, value]) => ({
					name,
					value: value as number
				}));

				if (joints.length > 0) {
					const command: RobotCommand = {
						timestamp: Date.now(),
						joints
					};
					this.notifyCommand(command);
				}
			});

			// Use workspace ID from config or default
			this.workspaceId = this.config.workspaceId || "default-workspace";

			let roomData;
			if (joinExistingRoom) {
				// Join existing room (for Inference Session integration)
				roomData = { workspaceId: this.workspaceId, roomId: this.config.robotId };
				console.log(
					`[RemoteConsumer] Joining existing room ${this.config.robotId} in workspace ${this.workspaceId}`
				);
			} else {
				// Create new room (for standalone operation)
				roomData = await this.client.createRoom(this.workspaceId, this.config.robotId);
				console.log(
					`[RemoteConsumer] Created new room ${roomData.roomId} in workspace ${roomData.workspaceId}`
				);
			}

			const success = await this.consumer.connect(roomData.workspaceId, roomData.roomId, this.id);

			if (!success) {
				throw new Error("Failed to connect consumer to room");
			}

			this._status = { isConnected: true, lastConnected: new Date() };
			this.notifyStatusChange();

			console.log(
				`[RemoteConsumer] Connected successfully to workspace ${roomData.workspaceId}, room ${roomData.roomId}`
			);
		} catch (error) {
			console.error(`[RemoteConsumer] Connection failed:`, error);
			this._status = { isConnected: false, error: `Connection failed: ${error}` };
			this.notifyStatusChange();
			throw error;
		}
	}

	async disconnect(): Promise<void> {
		console.log(`[RemoteConsumer] Disconnecting...`);

		if (this.consumer) {
			await this.consumer.disconnect();
			this.consumer = null;
		}
		if (this.client) {
			// Client doesn't need explicit disconnect
			this.client = null;
		}

		this.workspaceId = null;
		this._status = { isConnected: false };
		this.notifyStatusChange();

		console.log(`[RemoteConsumer] Disconnected`);
	}

	// Event handlers
	onStatusChange(callback: (status: ConnectionStatus) => void): () => void {
		this.statusCallbacks.push(callback);
		return () => {
			const index = this.statusCallbacks.indexOf(callback);
			if (index >= 0) {
				this.statusCallbacks.splice(index, 1);
			}
		};
	}

	onCommand(callback: (command: RobotCommand) => void): () => void {
		this.commandCallbacks.push(callback);
		return () => {
			const index = this.commandCallbacks.indexOf(callback);
			if (index >= 0) {
				this.commandCallbacks.splice(index, 1);
			}
		};
	}

	// Private methods
	private notifyCommand(command: RobotCommand): void {
		this.commandCallbacks.forEach((callback) => {
			try {
				callback(command);
			} catch (error) {
				console.error("[RemoteConsumer] Error in command callback:", error);
			}
		});
	}

	private notifyStatusChange(): void {
		this.statusCallbacks.forEach((callback) => {
			try {
				callback(this._status);
			} catch (error) {
				console.error("[RemoteConsumer] Error in status callback:", error);
			}
		});
	}
}