mecha10_core/

recorder.rs

//! Bag File Recording and Replay (ROS bag-like)
//!
//! Records messages from arbitrary topics to bag files and replays them later.
//! Similar to ROS bags but designed for the Mecha10 framework.
//!
//! # Features
//!
//! - Record messages from any topic with timestamps
//! - Replay messages with original or custom timing
//! - Multiple recording formats (CBOR, MessagePack, JSON)
//! - Compression support (zstd, gzip)
//! - Topic filtering on record and replay
//! - Time manipulation (speed up, slow down, step through)
//! - Metadata tracking (topics, message counts, duration)
//!
//! # File Format
//!
//! Bag files use a simple chunked format:
//!
//! ```text
//! [Header]
//! - Magic bytes: "METABAG1"
//! - Version: u32
//! - Compression: u8 (0=none, 1=zstd, 2=gzip)
//! - Metadata length: u64
//! - Metadata: CBOR-encoded BagMetadata
//!
//! [Chunk 0]
//! - Timestamp: u64 (microseconds since epoch)
//! - Topic length: u32
//! - Topic: String
//! - Message length: u32
//! - Message data: Vec<u8> (serialized message)
//!
//! [Chunk 1]
//! ...
//! ```
//!
//! # Example: Recording
//!
//! ```rust,no_run
//! use mecha10::prelude::*;
//! use mecha10::recorder::{BagRecorder, RecorderConfig};
//!
//! #[tokio::main]
//! async fn main() -> Result<()> {
//!     let ctx = Context::new("recorder").await?;
//!
//!     let config = RecorderConfig {
//!         output_path: "/tmp/recording.bag".into(),
//!         topics: vec![
//!             "/sensor/camera/rgb".to_string(),
//!             "/sensor/lidar/scan".to_string(),
//!         ],
//!         compression: Compression::Zstd,
//!         buffer_size: 1000,
//!     };
//!
//!     let mut recorder = BagRecorder::new(ctx, config).await?;
//!     recorder.start().await?;
//!
//!     // Record for 60 seconds
//!     tokio::time::sleep(Duration::from_secs(60)).await;
//!
//!     recorder.stop().await?;
//!     println!("Recording saved: {:?}", recorder.metadata());
//!
//!     Ok(())
//! }
//! ```
//!
//! # Example: Replay
//!
//! ```rust,no_run
//! use mecha10::prelude::*;
//! use mecha10::recorder::{BagPlayer, PlayerConfig};
//!
//! #[tokio::main]
//! async fn main() -> Result<()> {
//!     let ctx = Context::new("player").await?;
//!
//!     let config = PlayerConfig {
//!         bag_path: "/tmp/recording.bag".into(),
//!         topics: None, // Replay all topics
//!         rate: 1.0, // Real-time speed
//!         loop_playback: false,
//!         start_paused: false,
//!     };
//!
//!     let mut player = BagPlayer::new(ctx, config).await?;
//!     player.play().await?;
//!
//!     Ok(())
//! }
//! ```

use crate::{Context, Mecha10Error, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

// ============================================================================
// Constants
// ============================================================================

const MAGIC_BYTES: &[u8; 8] = b"METABAG1";
const VERSION: u32 = 1;

// ============================================================================
// Types
// ============================================================================

/// Compression format
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum Compression {
    /// No compression
    #[default]
    None = 0,
    /// Zstandard compression (fast, good ratio)
    Zstd = 1,
    /// Gzip compression (slower, better compression)
    Gzip = 2,
}

impl From<u8> for Compression {
    fn from(v: u8) -> Self {
        match v {
            1 => Compression::Zstd,
            2 => Compression::Gzip,
            _ => Compression::None,
        }
    }
}

/// Bag file metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BagMetadata {
    /// Topics recorded
    pub topics: Vec<String>,

    /// Message count per topic
    pub message_counts: HashMap<String, u64>,

    /// Start timestamp (microseconds since epoch)
    pub start_time_us: u64,

    /// End timestamp (microseconds since epoch)
    pub end_time_us: u64,

    /// Duration in seconds
    pub duration_sec: f64,

    /// Total message count
    pub total_messages: u64,

    /// Compression used
    pub compression: Compression,

    /// Custom metadata (key-value pairs)
    #[serde(default)]
    pub custom: HashMap<String, String>,
}

/// Message chunk in bag file
#[derive(Debug, Clone)]
struct MessageChunk {
    /// Timestamp (microseconds since epoch)
    timestamp_us: u64,

    /// Topic name
    topic: String,

    /// Serialized message data
    data: Vec<u8>,
}

// ============================================================================
// Recorder Configuration
// ============================================================================

/// Recorder configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecorderConfig {
    /// Output bag file path
    pub output_path: PathBuf,

    /// Topics to record (empty = record all)
    #[serde(default)]
    pub topics: Vec<String>,

    /// Compression format
    #[serde(default)]
    pub compression: Compression,

    /// Buffer size (number of messages to buffer before writing)
    #[serde(default = "default_buffer_size")]
    pub buffer_size: usize,

    /// Maximum file size in bytes (0 = unlimited)
    #[serde(default)]
    pub max_file_size: u64,

    /// Split files when max size reached
    #[serde(default)]
    pub split_on_max_size: bool,
}

fn default_buffer_size() -> usize {
    1000
}

impl Default for RecorderConfig {
    fn default() -> Self {
        Self {
            output_path: PathBuf::from("recording.bag"),
            topics: Vec::new(),
            compression: Compression::None,
            buffer_size: default_buffer_size(),
            max_file_size: 0,
            split_on_max_size: false,
        }
    }
}

// ============================================================================
// Bag Recorder
// ============================================================================

/// Bag file recorder service
///
/// Records messages from specified topics to a bag file.
pub struct BagRecorder {
    config: RecorderConfig,
    ctx: Arc<Context>,
    buffer: Arc<RwLock<Vec<MessageChunk>>>,
    metadata: Arc<RwLock<BagMetadata>>,
    running: Arc<RwLock<bool>>,
    write_task: Option<tokio::task::JoinHandle<()>>,
}

impl BagRecorder {
    /// Create a new bag recorder
    pub async fn new(ctx: Context, config: RecorderConfig) -> Result<Self> {
        let metadata = BagMetadata {
            topics: config.topics.clone(),
            message_counts: HashMap::new(),
            start_time_us: crate::prelude::now_micros(),
            end_time_us: 0,
            duration_sec: 0.0,
            total_messages: 0,
            compression: config.compression,
            custom: HashMap::new(),
        };

        Ok(Self {
            config,
            ctx: Arc::new(ctx),
            buffer: Arc::new(RwLock::new(Vec::new())),
            metadata: Arc::new(RwLock::new(metadata)),
            running: Arc::new(RwLock::new(false)),
            write_task: None,
        })
    }

    /// Start recording
    pub async fn start(&mut self) -> Result<()> {
        *self.running.write().await = true;

        info!("Starting bag recorder: {:?}", self.config.output_path);

        // Write header
        self.write_header().await?;

        // Spawn background task to write buffered messages
        let buffer = Arc::clone(&self.buffer);
        let _metadata = Arc::clone(&self.metadata);
        let running = Arc::clone(&self.running);
        let output_path = self.config.output_path.clone();
        let compression = self.config.compression;
        let buffer_size = self.config.buffer_size;

        let write_task = tokio::spawn(async move {
            while *running.read().await {
                tokio::time::sleep(Duration::from_millis(100)).await;

                let mut buf = buffer.write().await;
                if buf.len() >= buffer_size {
                    if let Err(e) = Self::write_chunks(&output_path, &buf, compression).await {
                        warn!("Failed to write chunks: {}", e);
                    }
                    buf.clear();
                }
            }

            // Final flush
            let buf = buffer.read().await;
            if !buf.is_empty() {
                if let Err(e) = Self::write_chunks(&output_path, &buf, compression).await {
                    warn!("Failed to write final chunks: {}", e);
                }
            }
        });

        self.write_task = Some(write_task);

        // Subscribe to topics and record messages
        for topic in &self.config.topics {
            let topic_clone = topic.clone();
            let buffer = Arc::clone(&self.buffer);
            let metadata = Arc::clone(&self.metadata);
            let running = Arc::clone(&self.running);
            let ctx = Arc::clone(&self.ctx);

            tokio::spawn(async move {
                // Subscribe to raw bytes to avoid deserialization
                match ctx.subscribe_raw::<serde_json::Value>(&topic_clone).await {
                    Ok(mut rx) => {
                        while *running.read().await {
                            if let Some(msg) = rx.recv().await {
                                // Serialize to CBOR
                                if let Ok(data) = serde_cbor::to_vec(&msg) {
                                    let chunk = MessageChunk {
                                        timestamp_us: crate::prelude::now_micros(),
                                        topic: topic_clone.clone(),
                                        data,
                                    };

                                    buffer.write().await.push(chunk);

                                    // Update metadata
                                    let mut meta = metadata.write().await;
                                    *meta.message_counts.entry(topic_clone.clone()).or_insert(0) += 1;
                                    meta.total_messages += 1;
                                }
                            }
                        }
                    }
                    Err(e) => {
                        warn!("Failed to subscribe to {}: {}", topic_clone, e);
                    }
                }
            });
        }

        Ok(())
    }

    /// Stop recording
    pub async fn stop(&mut self) -> Result<()> {
        *self.running.write().await = false;

        info!("Stopping bag recorder");

        // Wait for write task to complete
        if let Some(task) = self.write_task.take() {
            task.await
                .map_err(|e| Mecha10Error::Other(format!("Write task error: {}", e)))?;
        }

        // Update metadata
        let mut meta = self.metadata.write().await;
        meta.end_time_us = crate::prelude::now_micros();
        meta.duration_sec = (meta.end_time_us - meta.start_time_us) as f64 / 1_000_000.0;

        // Write final metadata (overwrite header)
        self.write_header().await?;

        info!(
            "Recording complete: {} messages, {:.1}s",
            meta.total_messages, meta.duration_sec
        );

        Ok(())
    }

    /// Get current metadata
    pub async fn metadata(&self) -> BagMetadata {
        self.metadata.read().await.clone()
    }

    /// Write bag file header with metadata
    async fn write_header(&self) -> Result<()> {
        let meta = self.metadata.read().await;

        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(&self.config.output_path)
            .map_err(Mecha10Error::Io)?;

        // Write magic bytes
        file.write_all(MAGIC_BYTES).map_err(Mecha10Error::Io)?;

        // Write version
        file.write_all(&VERSION.to_le_bytes()).map_err(Mecha10Error::Io)?;

        // Write compression
        file.write_all(&[self.config.compression as u8])
            .map_err(Mecha10Error::Io)?;

        // Write metadata
        let meta_bytes =
            serde_cbor::to_vec(&*meta).map_err(|e| Mecha10Error::Other(format!("Serialization error: {}", e)))?;

        file.write_all(&(meta_bytes.len() as u64).to_le_bytes())
            .map_err(Mecha10Error::Io)?;

        file.write_all(&meta_bytes).map_err(Mecha10Error::Io)?;

        Ok(())
    }

    /// Write message chunks to file
    async fn write_chunks(path: &PathBuf, chunks: &[MessageChunk], _compression: Compression) -> Result<()> {
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(path)
            .map_err(Mecha10Error::Io)?;

        for chunk in chunks {
            // Write timestamp
            file.write_all(&chunk.timestamp_us.to_le_bytes())
                .map_err(Mecha10Error::Io)?;

            // Write topic
            let topic_bytes = chunk.topic.as_bytes();
            file.write_all(&(topic_bytes.len() as u32).to_le_bytes())
                .map_err(Mecha10Error::Io)?;
            file.write_all(topic_bytes).map_err(Mecha10Error::Io)?;

            // Write message data
            file.write_all(&(chunk.data.len() as u32).to_le_bytes())
                .map_err(Mecha10Error::Io)?;
            file.write_all(&chunk.data).map_err(Mecha10Error::Io)?;
        }

        Ok(())
    }
}

// ============================================================================
// Player Configuration
// ============================================================================

/// Player configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PlayerConfig {
    /// Bag file path to play
    pub bag_path: PathBuf,

    /// Topics to replay (None = replay all)
    #[serde(default)]
    pub topics: Option<Vec<String>>,

    /// Playback rate (1.0 = real-time, 2.0 = 2x speed, 0.5 = half speed)
    #[serde(default = "default_rate")]
    pub rate: f64,

    /// Loop playback
    #[serde(default)]
    pub loop_playback: bool,

    /// Start paused (requires manual play command)
    #[serde(default)]
    pub start_paused: bool,

    /// Skip to timestamp (microseconds from start)
    #[serde(default)]
    pub skip_to_us: u64,
}

fn default_rate() -> f64 {
    1.0
}

impl Default for PlayerConfig {
    fn default() -> Self {
        Self {
            bag_path: PathBuf::from("recording.bag"),
            topics: None,
            rate: 1.0,
            loop_playback: false,
            start_paused: false,
            skip_to_us: 0,
        }
    }
}

// ============================================================================
// Bag Player
// ============================================================================

/// Bag file player service
///
/// Replays messages from a bag file to their original topics.
pub struct BagPlayer {
    config: PlayerConfig,
    ctx: Arc<Context>,
    metadata: BagMetadata,
    chunks: Vec<MessageChunk>,
    paused: Arc<RwLock<bool>>,
    playing: Arc<RwLock<bool>>,
}

impl BagPlayer {
    /// Create a new bag player
    pub async fn new(ctx: Context, config: PlayerConfig) -> Result<Self> {
        info!("Loading bag file: {:?}", config.bag_path);

        // Read and parse bag file
        let (metadata, chunks) = Self::read_bag_file(&config.bag_path).await?;

        info!(
            "Loaded bag: {} messages, {:.1}s duration, {} topics",
            metadata.total_messages,
            metadata.duration_sec,
            metadata.topics.len()
        );

        Ok(Self {
            config,
            ctx: Arc::new(ctx),
            metadata,
            chunks,
            paused: Arc::new(RwLock::new(false)),
            playing: Arc::new(RwLock::new(false)),
        })
    }

    /// Start playback
    pub async fn play(&mut self) -> Result<()> {
        if self.config.start_paused {
            *self.paused.write().await = true;
        }

        *self.playing.write().await = true;

        info!("Starting playback at {}x speed", self.config.rate);

        loop {
            let start_time = self.chunks.first().map(|c| c.timestamp_us).unwrap_or(0);
            let playback_start = std::time::Instant::now();

            for chunk in &self.chunks {
                // Wait for pause to be released
                while *self.paused.read().await {
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }

                if !*self.playing.read().await {
                    break;
                }

                // Apply topic filtering
                if let Some(topics) = &self.config.topics {
                    if !topics.contains(&chunk.topic) {
                        continue;
                    }
                }

                // Calculate sleep time based on playback rate
                let msg_time_offset = chunk.timestamp_us - start_time;
                let real_time_offset = (msg_time_offset as f64 / self.config.rate) as u64;
                let target_time = playback_start + Duration::from_micros(real_time_offset);

                if let Some(sleep_duration) = target_time.checked_duration_since(std::time::Instant::now()) {
                    tokio::time::sleep(sleep_duration).await;
                }

                // Deserialize and publish message
                match serde_cbor::from_slice::<serde_json::Value>(&chunk.data) {
                    Ok(msg) => {
                        if let Err(e) = self.ctx.publish_raw(&chunk.topic, &msg).await {
                            warn!("Failed to publish to {}: {}", chunk.topic, e);
                        } else {
                            debug!("Replayed message on {}", chunk.topic);
                        }
                    }
                    Err(e) => {
                        warn!("Failed to deserialize message: {}", e);
                    }
                }
            }

            if !self.config.loop_playback {
                break;
            }

            info!("Looping playback");
        }

        *self.playing.write().await = false;
        info!("Playback complete");

        Ok(())
    }

    /// Pause playback
    pub async fn pause(&self) {
        *self.paused.write().await = true;
        info!("Playback paused");
    }

    /// Resume playback
    pub async fn resume(&self) {
        *self.paused.write().await = false;
        info!("Playback resumed");
    }

    /// Stop playback
    pub async fn stop(&self) {
        *self.playing.write().await = false;
        info!("Playback stopped");
    }

    /// Get bag metadata
    pub fn metadata(&self) -> &BagMetadata {
        &self.metadata
    }

    /// Read and parse bag file
    async fn read_bag_file(path: &PathBuf) -> Result<(BagMetadata, Vec<MessageChunk>)> {
        let mut file = std::fs::File::open(path).map_err(Mecha10Error::Io)?;

        // Read magic bytes
        let mut magic = [0u8; 8];
        file.read_exact(&mut magic).map_err(Mecha10Error::Io)?;

        if &magic != MAGIC_BYTES {
            return Err(Mecha10Error::Other("Invalid bag file format".to_string()));
        }

        // Read version
        let mut version_bytes = [0u8; 4];
        file.read_exact(&mut version_bytes).map_err(Mecha10Error::Io)?;
        let version = u32::from_le_bytes(version_bytes);

        if version != VERSION {
            return Err(Mecha10Error::Other(format!("Unsupported bag version: {}", version)));
        }

        // Read compression
        let mut compression_byte = [0u8; 1];
        file.read_exact(&mut compression_byte).map_err(Mecha10Error::Io)?;
        let _compression = Compression::from(compression_byte[0]);

        // Read metadata
        let mut meta_len_bytes = [0u8; 8];
        file.read_exact(&mut meta_len_bytes).map_err(Mecha10Error::Io)?;
        let meta_len = u64::from_le_bytes(meta_len_bytes) as usize;

        let mut meta_bytes = vec![0u8; meta_len];
        file.read_exact(&mut meta_bytes).map_err(Mecha10Error::Io)?;

        let metadata: BagMetadata = serde_cbor::from_slice(&meta_bytes)
            .map_err(|e| Mecha10Error::Other(format!("Failed to parse metadata: {}", e)))?;

        // Read message chunks
        let mut chunks = Vec::new();

        loop {
            // Read timestamp
            let mut timestamp_bytes = [0u8; 8];
            if file.read_exact(&mut timestamp_bytes).is_err() {
                break; // End of file
            }
            let timestamp_us = u64::from_le_bytes(timestamp_bytes);

            // Read topic
            let mut topic_len_bytes = [0u8; 4];
            file.read_exact(&mut topic_len_bytes).map_err(Mecha10Error::Io)?;
            let topic_len = u32::from_le_bytes(topic_len_bytes) as usize;

            let mut topic_bytes = vec![0u8; topic_len];
            file.read_exact(&mut topic_bytes).map_err(Mecha10Error::Io)?;
            let topic = String::from_utf8(topic_bytes)
                .map_err(|e| Mecha10Error::Other(format!("Invalid topic string: {}", e)))?;

            // Read message data
            let mut data_len_bytes = [0u8; 4];
            file.read_exact(&mut data_len_bytes).map_err(Mecha10Error::Io)?;
            let data_len = u32::from_le_bytes(data_len_bytes) as usize;

            let mut data = vec![0u8; data_len];
            file.read_exact(&mut data).map_err(Mecha10Error::Io)?;

            chunks.push(MessageChunk {
                timestamp_us,
                topic,
                data,
            });
        }

        Ok((metadata, chunks))
    }
}