src/lib/robot/drivers/WebSocketSlave.ts

import type {
	SlaveDriver,
	DriverJointState,
	ConnectionStatus,
	RobotCommand,
	StateUpdateCallback,
	StatusChangeCallback,
	UnsubscribeFn
} from "$lib/types/robotDriver";

export interface WebSocketSlaveConfig {
	type: "websocket-slave";
	url: string;
	robotId: string;
	apiKey?: string;
}

/**
 * WebSocket Slave Driver
 * Connects to FastAPI WebSocket server as a slave to receive commands
 */
export class WebSocketSlave implements SlaveDriver {
	readonly type = "slave" as const;
	readonly id: string;
	readonly name: string;

	private _status: ConnectionStatus = { isConnected: false };
	private config: WebSocketSlaveConfig;

	// WebSocket connection
	private ws?: WebSocket;
	private reconnectAttempts = 0;
	private maxReconnectAttempts = 5;
	private reconnectDelay = 1000;

	// Joint states
	private jointStates: DriverJointState[] = [];

	// Event callbacks
	private stateCallbacks: StateUpdateCallback[] = [];
	private statusCallbacks: StatusChangeCallback[] = [];

	constructor(config: WebSocketSlaveConfig, initialJointStates: DriverJointState[]) {
		this.config = config;
		this.id = `websocket-slave-${Date.now()}`;
		this.name = `WebSocket Slave (${config.robotId})`;

		// Initialize joint states
		this.jointStates = initialJointStates.map((state) => ({ ...state }));

		console.log(
			`Created WebSocketSlave for robot ${config.robotId} with ${this.jointStates.length} joints`
		);
	}

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

	async connect(): Promise<void> {
		console.log(`Connecting ${this.name} to ${this.config.url}...`);

		try {
			// Build WebSocket URL
			const wsUrl = this.buildWebSocketUrl();

			// Create WebSocket connection
			this.ws = new WebSocket(wsUrl);

			// Set up event handlers
			this.setupWebSocketHandlers();

			// Wait for connection
			await this.waitForConnection();

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

			console.log(`${this.name} connected successfully`);
		} catch (error) {
			this._status = {
				isConnected: false,
				error: `Connection failed: ${error}`
			};
			this.notifyStatusChange();
			throw error;
		}
	}

	async disconnect(): Promise<void> {
		console.log(`Disconnecting ${this.name}...`);

		if (this.ws) {
			this.ws.close();
			this.ws = undefined;
		}

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

		console.log(`${this.name} disconnected`);
	}

	async executeCommand(command: RobotCommand): Promise<void> {
		if (!this._status.isConnected) {
			throw new Error("Cannot execute command: WebSocket slave not connected");
		}

		console.log(`WebSocketSlave executing command with ${command.joints.length} joint updates`);

		// Apply joint updates locally (for visualization)
		for (const jointUpdate of command.joints) {
			const joint = this.jointStates.find((j) => j.name === jointUpdate.name);
			if (joint) {
				joint.virtualValue = jointUpdate.value;
				joint.realValue = jointUpdate.value; // Simulate perfect execution
			}
		}

		// Send status update to server
		await this.sendStatusUpdate();

		// Notify state update
		this.notifyStateUpdate();
	}

	async executeCommands(commands: RobotCommand[]): Promise<void> {
		console.log(`WebSocketSlave executing batch of ${commands.length} commands`);

		for (const command of commands) {
			await this.executeCommand(command);

			// Small delay between commands
			if (commands.length > 1) {
				await new Promise((resolve) => setTimeout(resolve, 50));
			}
		}
	}

	async readJointStates(): Promise<DriverJointState[]> {
		if (!this._status.isConnected) {
			throw new Error("Cannot read states: WebSocket slave not connected");
		}

		return [...this.jointStates];
	}

	async writeJointState(jointName: string, value: number): Promise<void> {
		const command: RobotCommand = {
			timestamp: Date.now(),
			joints: [{ name: jointName, value }]
		};

		await this.executeCommand(command);
	}

	async writeJointStates(updates: { jointName: string; value: number }[]): Promise<void> {
		const command: RobotCommand = {
			timestamp: Date.now(),
			joints: updates.map((update) => ({ name: update.jointName, value: update.value }))
		};

		await this.executeCommand(command);
	}

	// Event subscription methods
	onStateUpdate(callback: StateUpdateCallback): UnsubscribeFn {
		this.stateCallbacks.push(callback);
		return () => {
			const index = this.stateCallbacks.indexOf(callback);
			if (index >= 0) {
				this.stateCallbacks.splice(index, 1);
			}
		};
	}

	onStatusChange(callback: StatusChangeCallback): UnsubscribeFn {
		this.statusCallbacks.push(callback);
		return () => {
			const index = this.statusCallbacks.indexOf(callback);
			if (index >= 0) {
				this.statusCallbacks.splice(index, 1);
			}
		};
	}

	// Private methods
	private buildWebSocketUrl(): string {
		const baseUrl = this.config.url.replace(/^http/, "ws");
		return `${baseUrl}/ws/slave/${this.config.robotId}`;
	}

	private setupWebSocketHandlers(): void {
		if (!this.ws) return;

		this.ws.onopen = () => {
			console.log(`WebSocket slave connected for robot ${this.config.robotId}`);
			this.reconnectAttempts = 0; // Reset on successful connection
		};

		this.ws.onmessage = (event) => {
			try {
				const message = JSON.parse(event.data);
				this.handleServerMessage(message);
			} catch (error) {
				console.error("Failed to parse WebSocket message:", error);
			}
		};

		this.ws.onclose = (event) => {
			console.log(
				`WebSocket slave closed for robot ${this.config.robotId}:`,
				event.code,
				event.reason
			);
			this.handleDisconnection();
		};

		this.ws.onerror = (error) => {
			console.error(`WebSocket slave error for robot ${this.config.robotId}:`, error);
			this._status = {
				isConnected: false,
				error: `WebSocket error: ${error}`
			};
			this.notifyStatusChange();
		};
	}

	private async waitForConnection(): Promise<void> {
		if (!this.ws) throw new Error("WebSocket not created");

		return new Promise((resolve, reject) => {
			const timeout = setTimeout(() => {
				reject(new Error("Connection timeout"));
			}, 10000); // 10 second timeout

			if (this.ws!.readyState === WebSocket.OPEN) {
				clearTimeout(timeout);
				resolve();
				return;
			}

			this.ws!.onopen = () => {
				clearTimeout(timeout);
				resolve();
			};

			this.ws!.onerror = (error) => {
				clearTimeout(timeout);
				reject(error);
			};
		});
	}

	private async sendMessage(message: unknown): Promise<void> {
		if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
			throw new Error("WebSocket not connected");
		}

		this.ws.send(JSON.stringify(message));
	}

	private async sendStatusUpdate(): Promise<void> {
		if (!this._status.isConnected) return;

		await this.sendMessage({
			type: "status_update",
			timestamp: new Date().toISOString(),
			data: {
				isConnected: this._status.isConnected,
				lastConnected: this._status.lastConnected?.toISOString(),
				error: this._status.error
			}
		});
	}

	private async sendJointStates(): Promise<void> {
		if (!this._status.isConnected) return;

		await this.sendMessage({
			type: "joint_states",
			timestamp: new Date().toISOString(),
			data: this.jointStates
		});
	}

	private handleServerMessage(message: unknown): void {
		if (typeof message !== "object" || message === null) return;

		const { type, data } = message as { type: string; data?: unknown };

		switch (type) {
			case "execute_command":
				if (data && typeof data === "object") {
					this.executeCommand(data as RobotCommand).catch((error) =>
						console.error("Failed to execute command from server:", error)
					);
				}
				break;

			case "execute_sequence":
				if (data && typeof data === "object") {
					const sequence = data as { commands: RobotCommand[] };
					this.executeCommands(sequence.commands).catch((error) =>
						console.error("Failed to execute sequence from server:", error)
					);
				}
				break;

			case "stop_sequence":
				console.log(`Stopping sequences on robot ${this.config.robotId}`);
				// For now, just log - in a real implementation, this would cancel ongoing sequences
				this.sendMessage({
					type: "status_update",
					timestamp: new Date().toISOString(),
					data: { message: "Sequences stopped", isConnected: true }
				}).catch((error) => console.error("Failed to send stop confirmation:", error));
				break;

			case "sync_state":
				console.log(`Received state sync for robot ${this.config.robotId}:`, data);
				break;

			default:
				console.warn(`Unknown message type from server: ${type}`);
		}
	}

	private handleDisconnection(): void {
		this._status = { isConnected: false };
		this.notifyStatusChange();

		// Attempt reconnection if not manually disconnected
		if (this.reconnectAttempts < this.maxReconnectAttempts) {
			this.attemptReconnection();
		}
	}

	private async attemptReconnection(): Promise<void> {
		this.reconnectAttempts++;
		const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1); // Exponential backoff

		console.log(
			`Attempting slave reconnection ${this.reconnectAttempts}/${this.maxReconnectAttempts} in ${delay}ms...`
		);

		setTimeout(async () => {
			try {
				await this.connect();
			} catch (error) {
				console.error(`Slave reconnection attempt ${this.reconnectAttempts} failed:`, error);
			}
		}, delay);
	}

	private notifyStateUpdate(): void {
		this.stateCallbacks.forEach((callback) => {
			try {
				callback([...this.jointStates]);
			} catch (error) {
				console.error("Error in state update callback:", error);
			}
		});
	}

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