src/lib/robot/Robot.svelte.ts

import type { RobotState } from "$lib/types/robot";
import type {
	MasterDriver,
	SlaveDriver,
	DriverJointState,
	ConnectionStatus,
	RobotCommand,
	CommandSequence,
	UnsubscribeFn
} from "$lib/types/robotDriver";
import type IUrdfJoint from "@/components/3d/robot/URDF/interfaces/IUrdfJoint";
import { getRobotPollingConfig } from "$lib/configs/performanceConfig";

/**
 * Enhanced joint state that combines URDF data with driver state
 */
export interface ManagedJointState {
	// URDF properties
	name: string;
	urdfJoint: IUrdfJoint;
	servoId?: number;

	// Current state
	virtualValue: number; // What the UI shows/controls
	realValue?: number; // What the physical robot reports
	commandedValue: number; // What was last commanded by master/manual
	isActive: boolean; // Whether this joint is being controlled

	// Calibration state
	calibrationOffset: number; // Offset to apply: realValue = rawValue + offset
	restPosition: number; // Expected position for this joint in rest mode

	// Sync state
	lastVirtualUpdate: Date;
	lastRealUpdate?: Date;
	lastCommandUpdate?: Date;
	syncError?: string;
}

/**
 * Calibration state for robot
 */
export interface CalibrationState {
	isCalibrated: boolean;
	calibrationTime?: Date;
	offsets: { [jointName: string]: number };
}

/**
 * Control source information
 */
export interface ControlState {
	hasActiveMaster: boolean;
	masterName?: string;
	manualControlEnabled: boolean;
	lastCommandSource: "master" | "manual" | "none";
	commandQueue: RobotCommand[];
}

/**
 * Robot instance that manages URDF visualization and master-slave connections
 */
export class Robot {
	readonly id: string;
	readonly robotState: RobotState;

	// Joint management
	private _joints = $state<ManagedJointState[]>([]);

	// Master driver (command source)
	private _master = $state<MasterDriver | undefined>(undefined);
	private _masterStatus = $state<ConnectionStatus>({ isConnected: false });

	// Slave drivers (execution targets)
	private _slaves = $state<SlaveDriver[]>([]);
	private _slaveStatuses = $state<Map<string, ConnectionStatus>>(new Map());

	// Control state
	private _controlState = $state<ControlState>({
		hasActiveMaster: false,
		manualControlEnabled: true,
		lastCommandSource: "none",
		commandQueue: []
	});

	// Calibration management
	private _calibrationState = $state<CalibrationState>({
		isCalibrated: false,
		offsets: {}
	});

	// Sync control - use performance config
	private syncIntervalId?: number;
	private readonly syncInterval = getRobotPollingConfig().ROBOT_SYNC_INTERVAL_MS;

	// Event subscriptions
	private unsubscribeCallbacks: UnsubscribeFn[] = [];

	// Reactive getters
	get joints() {
		return this._joints;
	}
	get master() {
		return this._master;
	}
	get masterStatus() {
		return this._masterStatus;
	}
	get slaves() {
		return this._slaves;
	}
	get slaveStatuses() {
		return this._slaveStatuses;
	}
	get controlState() {
		return this._controlState;
	}
	get calibrationState() {
		return this._calibrationState;
	}
	get isCalibrated() {
		return this._calibrationState.isCalibrated;
	}

	get activeJoints() {
		return this._joints.filter((joint) => joint.isActive);
	}

	get connectedSlaves() {
		return this._slaves.filter((slave) => this._slaveStatuses.get(slave.id)?.isConnected);
	}

	get manualControlEnabled() {
		return this._controlState.manualControlEnabled && !this._controlState.hasActiveMaster;
	}

	constructor(id: string, robotState: RobotState) {
		this.id = id;
		this.robotState = robotState;
		this.initializeJoints();
	}

	/**
	 * Initialize joint states from URDF
	 */
	private initializeJoints(): void {
		const servoMap = this.robotState.urdfConfig.jointNameIdMap || {};
		const restPositions = this.robotState.urdfConfig.restPosition || {};

		this._joints = this.robotState.robot.joints
			.filter((urdfJoint) => urdfJoint.type === "revolute" || urdfJoint.type === "continuous")
			.map((urdfJoint) => {
				const jointName = urdfJoint.name || "";
				const restPosition = restPositions[jointName] || 0;

				// Initialize URDF joint rotation to rest position
				if (urdfJoint.type === "revolute") {
					const radians = (restPosition * Math.PI) / 180;
					const axis = urdfJoint.axis_xyz || [0, 0, 1];
					urdfJoint.rotation = [radians * axis[0], radians * axis[1], radians * axis[2]];
				}

				return {
					name: jointName,
					urdfJoint,
					servoId: servoMap[jointName],
					virtualValue: restPosition,
					commandedValue: restPosition,
					realValue: undefined,
					isActive: !!servoMap[jointName],
					lastVirtualUpdate: new Date(),
					lastRealUpdate: undefined,
					lastCommandUpdate: undefined,
					syncError: undefined,
					calibrationOffset: 0,
					restPosition
				};
			});

		console.log(`Initialized ${this._joints.length} joints for robot ${this.id}`);
	}

	// ============= MASTER MANAGEMENT =============

	/**
	 * Connect a master driver (command source)
	 */
	async setMaster(master: MasterDriver): Promise<void> {
		// Disconnect existing master
		if (this._master) {
			await this.removeMaster();
		}

		this._master = master;
		this._masterStatus = master.status;

		// Subscribe to master events
		const statusUnsub = master.onStatusChange((status) => {
			this._masterStatus = status;
			this._controlState.hasActiveMaster = status.isConnected;
			this._controlState.manualControlEnabled = !status.isConnected;
		});

		const commandUnsub = master.onCommand((commands) => {
			this.handleMasterCommands(commands);
		});

		const sequenceUnsub = master.onSequence((sequence) => {
			this.handleMasterSequence(sequence);
		});

		this.unsubscribeCallbacks.push(statusUnsub, commandUnsub, sequenceUnsub);

		// Update control state
		this._controlState.hasActiveMaster = master.status.isConnected;
		this._controlState.manualControlEnabled = !master.status.isConnected;
		this._controlState.masterName = master.name;

		// Start the master (for masters that need to poll/generate commands)
		if (master.status.isConnected) {
			await master.start();
		}

		console.log(`Master ${master.name} connected to robot ${this.id}`);
	}

	/**
	 * Remove the current master
	 */
	async removeMaster(): Promise<void> {
		if (!this._master) return;

		if (this._master.status.isConnected) {
			// Stop the master first, then disconnect
			try {
				await this._master.stop();
			} catch (error) {
				console.warn(`Error stopping master ${this._master.name}:`, error);
			}

			await this._master.disconnect();
		}

		this._master = undefined;
		this._masterStatus = { isConnected: false };
		this._controlState.hasActiveMaster = false;
		this._controlState.manualControlEnabled = true;
		this._controlState.masterName = undefined;

		console.log(`Master removed from robot ${this.id}. Manual control restored.`);
	}

	// ============= SLAVE MANAGEMENT =============

	/**
	 * Add a slave driver (execution target)
	 */
	async addSlave(slave: SlaveDriver): Promise<void> {
		// Check if slave already exists
		if (this._slaves.find((s) => s.id === slave.id)) {
			throw new Error(`Slave ${slave.id} already connected to robot ${this.id}`);
		}

		// Connect the slave
		await slave.connect();

		this._slaves.push(slave);
		this._slaveStatuses.set(slave.id, slave.status);

		// Subscribe to slave events
		const statusUnsub = slave.onStatusChange((status) => {
			this._slaveStatuses.set(slave.id, status);
		});

		const stateUnsub = slave.onStateUpdate((driverStates) => {
			this.updateFromSlave(slave.id, driverStates);
		});

		this.unsubscribeCallbacks.push(statusUnsub, stateUnsub);

		// Move slave to rest position for safety
		await this.moveSlaveToRestPosition(slave);

		console.log(`Slave ${slave.name} (${slave.id}) connected to robot ${this.id}`);
	}

	/**
	 * Remove a slave driver
	 */
	async removeSlave(slaveId: string): Promise<void> {
		const slaveIndex = this._slaves.findIndex((s) => s.id === slaveId);
		if (slaveIndex === -1) return;

		const slave = this._slaves[slaveIndex];

		// Move to rest position before disconnection
		if (this._slaveStatuses.get(slaveId)?.isConnected) {
			try {
				await this.moveSlaveToRestPosition(slave);
				await new Promise((resolve) => setTimeout(resolve, 500));
			} catch (error) {
				console.warn(`Failed to move slave ${slaveId} to rest position:`, error);
			}
		}

		// Disconnect slave
		if (slave.status.isConnected) {
			await slave.disconnect();
		}

		// Remove from arrays
		this._slaves.splice(slaveIndex, 1);
		this._slaveStatuses.delete(slaveId);

		console.log(`Slave ${slaveId} removed from robot ${this.id}`);
	}

	// ============= COMMAND HANDLING =============

	/**
	 * Handle commands from master
	 */
	private async handleMasterCommands(commands: RobotCommand[]): Promise<void> {
		console.log(`Received ${commands.length} commands from master for robot ${this.id}`);

		// Update control state
		this._controlState.lastCommandSource = "master";
		this._controlState.commandQueue = [...commands];

		// Execute on all connected slaves
		const executePromises = this.connectedSlaves.map((slave) =>
			slave
				.executeCommands(commands)
				.catch((error) => console.error(`Slave ${slave.id} failed to execute commands:`, error))
		);

		await Promise.allSettled(executePromises);

		// Update virtual state from commands
		this.updateVirtualFromCommands(commands);
	}

	/**
	 * Handle command sequence from master
	 */
	private async handleMasterSequence(sequence: CommandSequence): Promise<void> {
		console.log(`Received sequence "${sequence.name}" from master for robot ${this.id}`);

		// Process sequence commands one by one
		for (const command of sequence.commands) {
			await this.handleMasterCommands([command]);

			// Wait for command duration if specified
			if (command.duration) {
				await new Promise((resolve) => setTimeout(resolve, command.duration));
			}
		}
	}

	/**
	 * Manual control - update joint value (when no master active)
	 */
	async updateJointValue(jointName: string, value: number): Promise<void> {
		if (!this.manualControlEnabled) {
			console.warn(`Manual control disabled - master is active for robot ${this.id}`);
			return;
		}

		const joint = this._joints.find((j) => j.name === jointName);
		if (!joint) {
			console.warn(`Joint ${jointName} not found in robot ${this.id}`);
			return;
		}

		// Apply limits
		const clampedValue = this.clampJointValue(joint, value);

		// Update virtual state
		joint.virtualValue = clampedValue;
		joint.commandedValue = clampedValue;
		joint.lastVirtualUpdate = new Date();
		joint.lastCommandUpdate = new Date();

		// Update URDF visualization
		this.updateUrdfJointRotation(joint, clampedValue);

		// Create command for slaves
		const command: RobotCommand = {
			timestamp: Date.now(),
			joints: [{ name: jointName, value: clampedValue }]
		};

		// Send to all connected slaves
		const executePromises = this.connectedSlaves.map((slave) =>
			slave
				.executeCommand(command)
				.catch((error) =>
					console.error(`Slave ${slave.id} failed to execute manual command:`, error)
				)
		);

		await Promise.allSettled(executePromises);

		this._controlState.lastCommandSource = "manual";
	}

	/**
	 * Manual control - update multiple joints
	 */
	async updateJointValues(updates: { jointName: string; value: number }[]): Promise<void> {
		if (!this.manualControlEnabled) {
			console.warn(`Manual control disabled - master is active for robot ${this.id}`);
			return;
		}

		// Update virtual states
		const commandJoints: { name: string; value: number }[] = [];

		for (const { jointName, value } of updates) {
			const joint = this._joints.find((j) => j.name === jointName);
			if (!joint) continue;

			const clampedValue = this.clampJointValue(joint, value);
			joint.virtualValue = clampedValue;
			joint.commandedValue = clampedValue;
			joint.lastVirtualUpdate = new Date();
			joint.lastCommandUpdate = new Date();

			// Update URDF visualization
			this.updateUrdfJointRotation(joint, clampedValue);

			commandJoints.push({ name: jointName, value: clampedValue });
		}

		if (commandJoints.length === 0) return;

		// Create batch command for slaves
		const command: RobotCommand = {
			timestamp: Date.now(),
			joints: commandJoints
		};

		// Send to all connected slaves
		const executePromises = this.connectedSlaves.map((slave) =>
			slave
				.executeCommand(command)
				.catch((error) =>
					console.error(`Slave ${slave.id} failed to execute manual batch command:`, error)
				)
		);

		await Promise.allSettled(executePromises);

		this._controlState.lastCommandSource = "manual";
	}

	// ============= STATE MANAGEMENT =============

	/**
	 * Update virtual state from master commands
	 */
	private updateVirtualFromCommands(commands: RobotCommand[]): void {
		for (const command of commands) {
			for (const jointCommand of command.joints) {
				const joint = this._joints.find((j) => j.name === jointCommand.name);
				if (joint) {
					joint.virtualValue = jointCommand.value;
					joint.commandedValue = jointCommand.value;
					joint.lastCommandUpdate = new Date();

					// Update URDF visualization
					this.updateUrdfJointRotation(joint, jointCommand.value);
				}
			}
		}
	}

	/**
	 * Update joint states from slave feedback
	 */
	private updateFromSlave(slaveId: string, driverStates: DriverJointState[]): void {
		for (const driverState of driverStates) {
			const joint = this._joints.find((j) => j.name === driverState.name);
			if (joint) {
				// Apply calibration offset to raw real value
				if (driverState.realValue !== undefined) {
					joint.realValue = driverState.realValue + joint.calibrationOffset;
				} else {
					joint.realValue = undefined;
				}

				joint.lastRealUpdate = new Date();
			}
		}
	}

	/**
	 * Update URDF joint rotation for 3D visualization
	 */
	private updateUrdfJointRotation(joint: ManagedJointState, degrees: number): void {
		const urdfJoint = joint.urdfJoint;
		if (!urdfJoint) return;

		if (urdfJoint.type === "revolute") {
			const radians = (degrees * Math.PI) / 180;
			const axis = urdfJoint.axis_xyz || [0, 0, 1];
			urdfJoint.rotation = [radians * axis[0], radians * axis[1], radians * axis[2]];
		}
	}

	/**
	 * Apply joint limits to a value
	 */
	private clampJointValue(joint: ManagedJointState, value: number): number {
		const limit = joint.urdfJoint.limit;
		if (!limit) return value;

		if (joint.urdfJoint.type === "revolute") {
			if (limit.lower !== undefined && limit.upper !== undefined) {
				const lowerDeg = (limit.lower * 180) / Math.PI;
				const upperDeg = (limit.upper * 180) / Math.PI;
				return Math.max(lowerDeg, Math.min(upperDeg, value));
			}
		} else if (joint.urdfJoint.type === "continuous") {
			if (limit.velocity !== undefined) {
				return Math.max(-limit.velocity, Math.min(limit.velocity, value));
			}
		}

		return value;
	}

	// ============= MOVEMENT FUNCTIONS =============

	/**
	 * Move robot to rest position smoothly
	 */
	async moveToRestPosition(durationMs: number = 3000): Promise<void> {
		console.log(`Moving robot ${this.id} to rest position over ${durationMs}ms`);

		const movements: Array<{
			joint: ManagedJointState;
			startValue: number;
			targetValue: number;
			totalDelta: number;
		}> = [];

		for (const joint of this._joints) {
			if (joint.isActive) {
				const startValue = joint.virtualValue;
				const targetValue = joint.restPosition;
				const totalDelta = targetValue - startValue;

				movements.push({
					joint,
					startValue,
					targetValue,
					totalDelta
				});
			}
		}

		if (movements.length === 0) {
			console.log("No active joints to move");
			return;
		}

		const startTime = Date.now();
		const updateInterval = 50;

		return new Promise((resolve) => {
			const animationStep = () => {
				const elapsed = Date.now() - startTime;
				const progress = Math.min(elapsed / durationMs, 1.0);

				const easedProgress =
					progress < 0.5
						? 4 * progress * progress * progress
						: 1 - Math.pow(-2 * progress + 2, 3) / 2;

				const updates: Array<{ jointName: string; value: number }> = [];

				for (const movement of movements) {
					const currentValue = movement.startValue + movement.totalDelta * easedProgress;
					updates.push({
						jointName: movement.joint.name,
						value: currentValue
					});
				}

				this.updateJointValues(updates);

				if (progress < 1.0) {
					setTimeout(animationStep, updateInterval);
				} else {
					console.log(`Robot ${this.id} reached rest position`);
					resolve();
				}
			};

			animationStep();
		});
	}

	/**
	 * Move a specific slave to rest position
	 */
	private async moveSlaveToRestPosition(slave: SlaveDriver): Promise<void> {
		const restCommand: RobotCommand = {
			timestamp: Date.now(),
			joints: this._joints.map((joint) => ({
				name: joint.name,
				value: joint.restPosition
			}))
		};

		await slave.executeCommand(restCommand);
	}

	// ============= CALIBRATION =============

	/**
	 * Calibrate robot using first connected slave
	 */
	async calibrateRobot(): Promise<void> {
		const connectedSlaves = this.connectedSlaves;
		if (connectedSlaves.length === 0) {
			throw new Error("Cannot calibrate: no slaves connected");
		}

		const primarySlave = connectedSlaves[0];
		console.log(`Calibrating robot ${this.id} using slave ${primarySlave.id}...`);

		try {
			const driverStates = await primarySlave.readJointStates();
			const offsets: { [jointName: string]: number } = {};

			for (const driverState of driverStates) {
				const joint = this._joints.find((j) => j.name === driverState.name);
				if (joint && driverState.realValue !== undefined) {
					const offset = joint.restPosition - driverState.realValue;
					offsets[joint.name] = offset;
					joint.calibrationOffset = offset;

					console.log(
						`Joint ${joint.name}: rest=${joint.restPosition}° real=${driverState.realValue}° offset=${offset.toFixed(1)}°`
					);
				}
			}

			this._calibrationState = {
				isCalibrated: true,
				calibrationTime: new Date(),
				offsets
			};

			console.log(`Robot ${this.id} calibrated successfully`);
		} catch (error) {
			console.error(`Calibration failed for robot ${this.id}:`, error);
			throw error;
		}
	}

	/**
	 * Clear calibration
	 */
	clearCalibration(): void {
		this._calibrationState = {
			isCalibrated: false,
			offsets: {}
		};

		this._joints.forEach((joint) => {
			joint.calibrationOffset = 0;
		});

		console.log(`Cleared calibration for robot ${this.id}`);
	}

	// ============= LIFECYCLE =============

	/**
	 * Get joint by name
	 */
	getJoint(name: string): ManagedJointState | undefined {
		return this._joints.find((j) => j.name === name);
	}

	/**
	 * Clean up resources
	 */
	async destroy(): Promise<void> {
		// Clean up subscriptions
		this.unsubscribeCallbacks.forEach((unsub) => unsub());
		this.unsubscribeCallbacks = [];

		// Remove master
		await this.removeMaster();

		// Remove all slaves
		const slaveIds = this._slaves.map((s) => s.id);
		for (const slaveId of slaveIds) {
			await this.removeSlave(slaveId);
		}

		console.log(`Robot ${this.id} destroyed`);
	}
}