import React, { useState, useEffect } from "react"; import { useNavigate, useLocation } from "react-router-dom"; import { Button } from "@/components/ui/button"; import { useToast } from "@/hooks/use-toast"; import { ArrowLeft, Square, SkipForward, RotateCcw, Play } from "lucide-react"; import UrdfViewer from "@/components/UrdfViewer"; import UrdfProcessorInitializer from "@/components/UrdfProcessorInitializer"; import { useApi } from "@/contexts/ApiContext"; interface RecordingConfig { leader_port: string; follower_port: string; leader_config: string; follower_config: string; dataset_repo_id: string; single_task: string; num_episodes: number; episode_time_s: number; reset_time_s: number; fps: number; video: boolean; push_to_hub: boolean; resume: boolean; } interface BackendStatus { recording_active: boolean; current_phase: string; current_episode?: number; total_episodes?: number; saved_episodes?: number; phase_elapsed_seconds?: number; phase_time_limit_s?: number; session_elapsed_seconds?: number; session_ended?: boolean; available_controls: { stop_recording: boolean; exit_early: boolean; rerecord_episode: boolean; }; } const Recording = () => { const location = useLocation(); const navigate = useNavigate(); const { toast } = useToast(); const { baseUrl, wsBaseUrl, fetchWithHeaders } = useApi(); // Get recording config from navigation state const recordingConfig = location.state?.recordingConfig as RecordingConfig; // Backend status state - this is the single source of truth const [backendStatus, setBackendStatus] = useState( null ); const [recordingSessionStarted, setRecordingSessionStarted] = useState(false); // Local UI state for immediate user feedback const [transitioningToReset, setTransitioningToReset] = useState(false); const [transitioningToNext, setTransitioningToNext] = useState(false); // Redirect if no config provided useEffect(() => { if (!recordingConfig) { toast({ title: "No Configuration", description: "Please start recording from the main page.", variant: "destructive", }); navigate("/"); } }, [recordingConfig, navigate, toast]); // Start recording session when component loads useEffect(() => { if (recordingConfig && !recordingSessionStarted) { startRecordingSession(); } }, [recordingConfig, recordingSessionStarted]); // Poll backend status continuously to stay in sync useEffect(() => { let statusInterval: NodeJS.Timeout; if (recordingSessionStarted) { const pollStatus = async () => { try { const response = await fetchWithHeaders( `${baseUrl}/recording-status` ); if (response.ok) { const status = await response.json(); console.log( `📊 Backend Status: ${status.current_phase} | Transition States: reset=${transitioningToReset}, next=${transitioningToNext}` ); setBackendStatus(status); // 🎯 CLEAR TRANSITION STATES: Only clear when backend actually reaches the expected phase if (status.current_phase === "resetting" && transitioningToReset) { console.log( "✅ Clearing transitioningToReset - backend reached resetting phase" ); setTransitioningToReset(false); } if (status.current_phase === "recording" && transitioningToNext) { console.log( "✅ Clearing transitioningToNext - backend reached recording phase" ); setTransitioningToNext(false); } // If backend recording stopped and session ended, navigate to upload if ( !status.recording_active && status.session_ended && recordingSessionStarted ) { // Navigate to upload window with dataset info const datasetInfo = { dataset_repo_id: recordingConfig.dataset_repo_id, single_task: recordingConfig.single_task, num_episodes: recordingConfig.num_episodes, saved_episodes: status.saved_episodes || 0, session_elapsed_seconds: status.session_elapsed_seconds || 0, }; navigate("/upload", { state: { datasetInfo } }); return; // Stop polling after navigation } } } catch (error) { console.error("Error polling recording status:", error); } }; // Poll immediately and then every second for real-time updates pollStatus(); statusInterval = setInterval(pollStatus, 1000); } return () => { if (statusInterval) clearInterval(statusInterval); }; }, [ recordingSessionStarted, recordingConfig, navigate, toast, transitioningToReset, transitioningToNext, ]); const formatTime = (seconds: number): string => { const mins = Math.floor(seconds / 60); const secs = seconds % 60; return `${mins.toString().padStart(2, "0")}:${secs .toString() .padStart(2, "0")}`; }; const startRecordingSession = async () => { try { const response = await fetchWithHeaders(`${baseUrl}/start-recording`, { method: "POST", body: JSON.stringify(recordingConfig), }); const data = await response.json(); if (response.ok) { setRecordingSessionStarted(true); toast({ title: "Recording Started", description: `Started recording ${recordingConfig.num_episodes} episodes`, }); } else { toast({ title: "Error Starting Recording", description: data.message || "Failed to start recording session.", variant: "destructive", }); navigate("/"); } } catch (error) { toast({ title: "Connection Error", description: "Could not connect to the backend server.", variant: "destructive", }); navigate("/"); } }; // Equivalent to pressing RIGHT ARROW key in original record.py const handleExitEarly = async () => { if (!backendStatus?.available_controls.exit_early) return; // 🎯 IMMEDIATE UI FEEDBACK: Show transition state before backend response const currentPhase = backendStatus.current_phase; if (currentPhase === "recording") { console.log("🎯 Setting transitioningToReset = true"); setTransitioningToReset(true); toast({ title: "Ending Episode Recording", description: `Moving to reset phase for episode ${backendStatus.current_episode}...`, }); } else if (currentPhase === "resetting") { console.log("🎯 Setting transitioningToNext = true"); setTransitioningToNext(true); toast({ title: "Reset Complete", description: `Moving to next episode...`, }); } try { const response = await fetchWithHeaders( `${baseUrl}/recording-exit-early`, { method: "POST", } ); const data = await response.json(); if (response.ok) { // ✅ SUCCESS: Don't clear transition states here - let them persist until backend phase changes // The transition states will be cleared when the backend status actually updates to the new phase } else { // Clear transition states on error setTransitioningToReset(false); setTransitioningToNext(false); toast({ title: "Error", description: data.message, variant: "destructive", }); } } catch (error) { // Clear transition states on error setTransitioningToReset(false); setTransitioningToNext(false); toast({ title: "Connection Error", description: "Could not connect to the backend server.", variant: "destructive", }); } }; // Equivalent to pressing LEFT ARROW key in original record.py const handleRerecordEpisode = async () => { if (!backendStatus?.available_controls.rerecord_episode) return; try { const response = await fetchWithHeaders( `${baseUrl}/recording-rerecord-episode`, { method: "POST", } ); const data = await response.json(); if (response.ok) { toast({ title: "Re-recording Episode", description: `Episode ${backendStatus.current_episode} will be re-recorded.`, }); } else { toast({ title: "Error", description: data.message, variant: "destructive", }); } } catch (error) { toast({ title: "Connection Error", description: "Could not connect to the backend server.", variant: "destructive", }); } }; // Equivalent to pressing ESC key in original record.py const handleStopRecording = async () => { try { const response = await fetchWithHeaders(`${baseUrl}/stop-recording`, { method: "POST", }); toast({ title: "Recording Stopped", description: "Recording session has been stopped.", }); // Navigate to upload window with current dataset info const datasetInfo = { dataset_repo_id: recordingConfig.dataset_repo_id, single_task: recordingConfig.single_task, num_episodes: recordingConfig.num_episodes, saved_episodes: backendStatus?.saved_episodes || 0, session_elapsed_seconds: backendStatus?.session_elapsed_seconds || 0, }; navigate("/upload", { state: { datasetInfo } }); } catch (error) { toast({ title: "Error", description: "Failed to stop recording.", variant: "destructive", }); } }; if (!recordingConfig) { return (

No recording configuration found.

); } // Show loading state while waiting for backend status if (!backendStatus) { return (

Connecting to recording session...

); } const currentPhase = backendStatus.current_phase; const currentEpisode = backendStatus.current_episode || 1; const totalEpisodes = backendStatus.total_episodes || recordingConfig.num_episodes; const phaseElapsedTime = backendStatus.phase_elapsed_seconds || 0; const phaseTimeLimit = backendStatus.phase_time_limit_s || (currentPhase === "recording" ? recordingConfig.episode_time_s : recordingConfig.reset_time_s); const sessionElapsedTime = backendStatus.session_elapsed_seconds || 0; const getPhaseTitle = () => { // 🎯 IMMEDIATE FEEDBACK: Show transition titles if (transitioningToReset) return "Transitioning to Reset"; if (transitioningToNext) return "Moving to Next Episode"; if (currentPhase === "recording") return "Episode Recording Time"; if (currentPhase === "resetting") return "Environment Reset Time"; return "Phase Time"; }; const getStatusText = () => { // 🎯 IMMEDIATE FEEDBACK: Show transition states if (transitioningToReset) return "MOVING TO RESET PHASE"; if (transitioningToNext) return "MOVING TO NEXT EPISODE"; if (currentPhase === "recording") return `RECORDING EPISODE ${currentEpisode}`; if (currentPhase === "resetting") return "RESET THE ENVIRONMENT"; if (currentPhase === "preparing") return "PREPARING SESSION"; return "SESSION COMPLETE"; }; const getStatusColor = () => { // 🎯 IMMEDIATE FEEDBACK: Show transition state colors if (transitioningToReset) return "text-blue-400"; // Blue for transition if (transitioningToNext) return "text-blue-400"; // Blue for transition if (currentPhase === "recording") return "text-red-400"; if (currentPhase === "resetting") return "text-orange-400"; if (currentPhase === "preparing") return "text-yellow-400"; return "text-gray-400"; }; const getDotColor = () => { // 🎯 IMMEDIATE FEEDBACK: Show transition state dots with animation if (transitioningToReset) return "bg-blue-500 animate-pulse"; // Blue pulsing for transition if (transitioningToNext) return "bg-blue-500 animate-pulse"; // Blue pulsing for transition if (currentPhase === "recording") return "bg-red-500 animate-pulse"; if (currentPhase === "resetting") return "bg-orange-500 animate-pulse"; if (currentPhase === "preparing") return "bg-yellow-500"; return "bg-gray-500"; }; return (
{/* Header */}

Recording Session

{/* Main Recording Dashboard */}
{/* Phase Timer */}

{getPhaseTitle()}

{formatTime(phaseElapsedTime)}
/ {formatTime(phaseTimeLimit)}
{/* Episode Progress */}

Episode Progress

{currentEpisode} of {totalEpisodes}
{recordingConfig.single_task}
{/* Session Timer */}

Total Session Time

{formatTime(sessionElapsedTime)}
Dataset: {recordingConfig.dataset_repo_id}
{/* Status and Controls */}
{/* Recording Status - takes up 3 columns */}
{/* Status header */}

Recording Status

{getStatusText()}
{/* Recording Phase Controls */} {currentPhase === "recording" && (
)} {/* Reset Phase Controls */} {currentPhase === "resetting" && (
)} {currentPhase === "completed" && (

✅ Recording session completed successfully!

Dataset:{" "} {recordingConfig.dataset_repo_id}

You will be redirected to the upload window shortly...

)} {/* Instructions */}

{currentPhase === "recording" ? "Episode Recording Instructions:" : currentPhase === "resetting" ? "Environment Reset Instructions:" : "Session Instructions:"}

{currentPhase === "recording" && (
  • • End Episode: Complete current episode and enter reset phase (Right Arrow)
  • • Re-record Episode: Restart current episode after reset phase (Left Arrow)
  • • Auto-end: Episode ends automatically after {formatTime(phaseTimeLimit)}
  • • Stop Recording: End entire session (ESC key)
)} {currentPhase === "resetting" && (
  • • Continue to Next Phase: Skip reset phase and continue (Right Arrow)
  • • Auto-continue: Automatically continues after {formatTime(phaseTimeLimit)}
  • • Reset Phase: Use this time to prepare your environment for the next episode
  • • Stop Recording: End entire session (ESC key)
)}
{/* URDF Viewer Section */}

Robot Visualizer

); }; export default Recording;