burn_telemetry 0.1.0

//! Rerun metric logger implementation

use super::RerunConfig;
use crate::metric::numeric_value;
use burn::train::{
    logger::MetricLogger,
    metric::{MetricEntry, NumericEntry},
};
use rerun::{RecordingStream, RecordingStreamBuilder};
use std::{
    collections::{HashMap, HashSet},
    fs,
    path::PathBuf,
    str::FromStr,
    sync::{Arc, Mutex, Weak},
};

/// Video stream state for tracking frame indices
#[derive(Debug, Clone)]
struct VideoStreamState {
    frame_index: i64,
    entity_path: String,
}

/// Handle to a video stream
///
/// When all handles to a stream are dropped, the stream is automatically cleaned up.
/// This prevents resource leaks and provides a type-safe API for stream management.
#[derive(Debug, Clone)]
pub struct VideoStreamHandle {
    key: String,
    _cleanup: Arc<StreamCleanup>,
}

/// Helper struct to track when all handles are dropped
#[derive(Debug)]
struct StreamCleanup {
    key: String,
    logger: Weak<Mutex<HashMap<String, VideoStreamState>>>,
}

impl Drop for StreamCleanup {
    fn drop(&mut self) {
        // Automatically cleanup stream when last handle is dropped
        if let Some(logger) = self.logger.upgrade() {
            if let Ok(mut streams) = logger.lock() {
                streams.remove(&self.key);
            }
        }
    }
}

/// Rerun metric logger for Burn training
///
/// This logger exports training metrics to Rerun for rich visualization
/// including support for images, videos, tensors, and other visual formats.
pub struct RerunLogger {
    rec: RecordingStream,
    current_epoch: usize,
    timeline_name: String,
    video_streams: Arc<Mutex<HashMap<String, VideoStreamState>>>,
}

impl RerunLogger {
    /// Create a new Rerun logger
    ///
    /// # Arguments
    ///
    /// * `config` - Configuration for the logger
    /// * `job_name` - Name of the training job
    ///
    /// # Returns
    ///
    /// A new logger instance or an error if initialization fails
    pub fn new(config: RerunConfig, job_name: impl Into<String>) -> Result<Self, String> {
        let job_name = job_name.into();
        let RerunConfig {
            application_id,
            recording_id,
            spawn_viewer,
            connect_addr,
            save_path,
            record_to_path,
            default_timeline,
        } = config;

        // Create recording stream builder
        let mut builder = RecordingStreamBuilder::new(application_id);

        if let Some(recording_id) = recording_id {
            builder = builder.recording_id(recording_id);
        }

        let rec = Self::initialize_stream(
            builder,
            spawn_viewer,
            connect_addr,
            save_path,
            record_to_path,
        )?;

        // Set application-level metadata
        rec.set_time_sequence("job", 0);
        rec.log_static(
            "world/job_info",
            &rerun::TextDocument::new(format!("Training job: {}", job_name)),
        )
        .ok();

        Ok(Self {
            rec,
            current_epoch: 1,
            timeline_name: default_timeline,
            video_streams: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    fn initialize_stream(
        builder: RecordingStreamBuilder,
        spawn_viewer: bool,
        connect_addr: Option<String>,
        save_path: Option<String>,
        record_to_path: Option<PathBuf>,
    ) -> Result<RecordingStream, String> {
        if let Some(sinks) = Self::build_sinks(
            spawn_viewer,
            connect_addr.as_deref(),
            save_path.as_deref(),
            record_to_path,
        )? {
            return builder
                .set_sinks(sinks)
                .map_err(|e| format!("Failed to initialize Rerun stream: {e}"));
        }

        let (rec, _storage) = builder
            .memory()
            .map_err(|e| format!("Failed to create buffered Rerun stream: {e}"))?;
        Ok(rec)
    }

    fn build_sinks(
        spawn_viewer: bool,
        connect_addr: Option<&str>,
        save_path: Option<&str>,
        record_to_path: Option<PathBuf>,
    ) -> Result<Option<Vec<Box<dyn rerun::sink::LogSink>>>, String> {
        let mut sinks: Vec<Box<dyn rerun::sink::LogSink>> = Vec::new();
        let mut seen_paths = HashSet::new();

        if spawn_viewer {
            sinks.push(Box::new(Self::spawn_viewer_sink()?));
        } else if let Some(addr) = connect_addr {
            sinks.push(Box::new(Self::grpc_sink_for_endpoint(addr)?));
        }

        if let Some(path) = save_path {
            Self::add_file_sink(&mut sinks, &mut seen_paths, PathBuf::from(path))?;
        }

        if let Some(path) = record_to_path {
            Self::add_file_sink(&mut sinks, &mut seen_paths, path)?;
        }

        if sinks.is_empty() {
            Ok(None)
        } else {
            Ok(Some(sinks))
        }
    }

    fn spawn_viewer_sink() -> Result<rerun::sink::GrpcSink, String> {
        let opts = rerun::SpawnOptions::default();
        rerun::spawn(&opts).map_err(|e| format!("Failed to spawn Rerun viewer: {e}"))?;
        let url = format!("rerun+http://{}/proxy", opts.connect_addr());
        Self::grpc_sink_from_url(&url)
    }

    fn grpc_sink_for_endpoint(endpoint: &str) -> Result<rerun::sink::GrpcSink, String> {
        let url = if endpoint.contains("://") {
            endpoint.to_owned()
        } else {
            format!("rerun+http://{endpoint}/proxy")
        };
        Self::grpc_sink_from_url(&url)
    }

    fn grpc_sink_from_url(url: &str) -> Result<rerun::sink::GrpcSink, String> {
        let uri = rerun::external::re_uri::ProxyUri::from_str(url)
            .map_err(|e| format!("Failed to parse Rerun endpoint {url}: {e}"))?;
        Ok(rerun::sink::GrpcSink::new(uri))
    }

    fn add_file_sink(
        sinks: &mut Vec<Box<dyn rerun::sink::LogSink>>,
        seen_paths: &mut HashSet<PathBuf>,
        path: PathBuf,
    ) -> Result<(), String> {
        if !seen_paths.insert(path.clone()) {
            return Ok(());
        }

        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)
                .map_err(|e| format!("Failed to create directory {}: {}", parent.display(), e))?;
        }

        let sink = rerun::sink::FileSink::new(path.clone()).map_err(|e| {
            format!(
                "Failed to create Rerun recording at {}: {}",
                path.display(),
                e
            )
        })?;
        sinks.push(Box::new(sink));
        Ok(())
    }

    /// Get the recording stream for custom logging
    pub fn recording_stream(&self) -> &RecordingStream {
        &self.rec
    }

    /// Log a scalar metric
    pub fn log_scalar(&self, name: &str, value: f64, epoch: usize) {
        self.rec
            .set_time_sequence(self.timeline_name.as_str(), epoch as i64);
        let path = resolve_entity_path(name, "metrics");
        self.rec.log(path, &rerun::Scalars::single(value)).ok();
    }

    /// Log a text entry
    pub fn log_text(&self, name: &str, text: &str, epoch: usize) {
        self.rec
            .set_time_sequence(self.timeline_name.as_str(), epoch as i64);
        let path = resolve_entity_path(name, "annotations");
        self.rec.log(path, &rerun::TextDocument::new(text)).ok();
    }

    /// Log an image (requires a Rerun Image)
    pub fn log_image(&self, name: &str, image: rerun::Image, epoch: usize) {
        self.rec
            .set_time_sequence(self.timeline_name.as_str(), epoch as i64);
        let path = resolve_entity_path(name, "images");
        self.rec.log(path, &image).ok();
    }

    /// Log a tensor (requires a Rerun Tensor)
    pub fn log_tensor(&self, name: &str, tensor: rerun::Tensor, step: usize) {
        self.rec
            .set_time_sequence(self.timeline_name.as_str(), step as i64);
        let path = resolve_entity_path(name, "tensors");
        self.rec.log(path, &tensor).ok();
    }

    /// Initialize or get a video stream for streaming video frames
    ///
    /// This creates a persistent video stream that can receive frames incrementally.
    /// Each stream maintains its own frame counter. The stream is automatically cleaned up
    /// when all handles are dropped.
    ///
    /// # Arguments
    ///
    /// * `stream_name` - Unique name for the video stream
    ///
    /// # Returns
    ///
    /// A handle to the video stream
    pub fn init_video_stream(&self, stream_name: impl Into<String>) -> VideoStreamHandle {
        let stream_name = stream_name.into();
        let stream_key = format!("video_stream_{}", stream_name.replace(' ', "_"));

        let mut streams = self.video_streams.lock().unwrap();
        if !streams.contains_key(&stream_key) {
            let entity_path = format!("video_streams/{}", stream_name.replace(' ', "_"));
            streams.insert(
                stream_key.clone(),
                VideoStreamState {
                    frame_index: 0,
                    entity_path,
                },
            );
        }
        drop(streams); // Release lock before creating cleanup

        VideoStreamHandle {
            key: stream_key.clone(),
            _cleanup: Arc::new(StreamCleanup {
                key: stream_key,
                logger: Arc::downgrade(&self.video_streams),
            }),
        }
    }

    /// Log a single frame to a video stream
    ///
    /// This method logs frames incrementally to a persistent video stream.
    /// Frames are automatically indexed and can be played back in the Rerun viewer.
    ///
    /// # Arguments
    ///
    /// * `handle` - The stream handle returned from `init_video_stream`
    /// * `frame` - The image frame to log
    /// * `epoch` - Current training epoch
    ///
    /// # Example
    ///
    /// ```ignore
    /// let stream_handle = logger.init_video_stream("training_progress");
    /// for frame in frames {
    ///     logger.log_video_frame(&stream_handle, frame, epoch);
    /// }
    /// ```
    pub fn log_video_frame(&self, handle: &VideoStreamHandle, frame: rerun::Image, epoch: usize) {
        let mut streams = self.video_streams.lock().unwrap();
        if let Some(state) = streams.get_mut(&handle.key) {
            let epoch_i64 = epoch as i64;
            let next_frame_index = state.frame_index + 1;
            let timeline_value = epoch_i64.max(next_frame_index);
            self.rec
                .set_time_sequence(self.timeline_name.as_str(), timeline_value);
            self.rec.set_time_sequence("frame", state.frame_index);

            self.rec.log(state.entity_path.as_str(), &frame).ok();
            state.frame_index = next_frame_index;
        }
    }

    /// Reset a video stream's frame counter
    ///
    /// This resets the frame index back to 0 for a stream, useful when
    /// starting a new video sequence in the same stream.
    ///
    /// # Arguments
    ///
    /// * `handle` - The stream handle to reset
    pub fn reset_video_stream(&self, handle: &VideoStreamHandle) {
        let mut streams = self.video_streams.lock().unwrap();
        if let Some(state) = streams.get_mut(&handle.key) {
            state.frame_index = 0;
        }
    }

    /// Get the current frame index for a video stream
    ///
    /// # Arguments
    ///
    /// * `handle` - The stream handle to query
    ///
    /// # Returns
    ///
    /// The current frame index, or None if the stream doesn't exist
    pub fn get_stream_frame_index(&self, handle: &VideoStreamHandle) -> Option<i64> {
        let streams = self.video_streams.lock().unwrap();
        streams.get(&handle.key).map(|s| s.frame_index)
    }
}

fn resolve_entity_path(name: &str, default_root: &str) -> String {
    let trimmed = name.trim();
    if trimmed.is_empty() {
        return default_root.to_string();
    }

    let has_path = trimmed.contains(['/', '\\']);
    let sanitized = sanitize_path(trimmed);

    if has_path {
        sanitized
    } else {
        format!("{}/{}", default_root, sanitized)
    }
}

fn sanitize_path(path: &str) -> String {
    path.split(['/', '\\'])
        .filter(|segment| !segment.is_empty())
        .map(sanitize_segment)
        .collect::<Vec<_>>()
        .join("/")
}

fn sanitize_segment(segment: &str) -> String {
    let mut cleaned = String::with_capacity(segment.len());
    for ch in segment.chars() {
        if ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '.') {
            cleaned.push(ch);
        } else {
            cleaned.push('_');
        }
    }
    if cleaned.is_empty() {
        cleaned.push_str("unnamed");
    }
    cleaned
}

impl MetricLogger for RerunLogger {
    fn log(&mut self, item: &MetricEntry) {
        if let Some(value) = numeric_value(item) {
            self.log_scalar(&item.name, value, self.current_epoch);
        } else {
            // Log as text if not numeric
            self.log_text(&item.name, &item.formatted, self.current_epoch);
        }
    }

    fn end_epoch(&mut self, epoch: usize) {
        self.current_epoch = epoch + 1;
    }

    fn read_numeric(&mut self, _name: &str, _epoch: usize) -> Result<Vec<NumericEntry>, String> {
        // Rerun is write-only; we don't support reading back
        Err("Rerun logger does not support reading metrics".to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_logger_creation() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job");
        // Buffered mode should work without viewer
        assert!(logger.is_ok());
    }

    #[test]
    fn test_read_not_supported() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let mut logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let result = logger.read_numeric("test", 1);
        assert!(result.is_err());
    }

    #[test]
    fn test_end_epoch() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let mut logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        assert_eq!(logger.current_epoch, 1);
        logger.end_epoch(1);
        assert_eq!(logger.current_epoch, 2);
        logger.end_epoch(2);
        assert_eq!(logger.current_epoch, 3);
    }

    #[test]
    fn test_init_video_stream() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let stream_handle = logger.init_video_stream("test_stream");
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(0));
    }

    #[test]
    fn test_video_stream_frame_counter() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let stream_handle = logger.init_video_stream("test_stream");

        // Create a simple test image
        let img = rerun::Image::from_color_model_and_bytes(
            vec![0u8; 8 * 8 * 3],
            [8, 8],
            rerun::ColorModel::RGB,
            rerun::ChannelDatatype::U8,
        );

        // Log several frames
        logger.log_video_frame(&stream_handle, img.clone(), 1);
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(1));

        logger.log_video_frame(&stream_handle, img.clone(), 1);
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(2));

        logger.log_video_frame(&stream_handle, img, 1);
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(3));
    }

    #[test]
    fn test_reset_video_stream() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let stream_handle = logger.init_video_stream("test_stream");

        let img = rerun::Image::from_color_model_and_bytes(
            vec![0u8; 8 * 8 * 3],
            [8, 8],
            rerun::ColorModel::RGB,
            rerun::ChannelDatatype::U8,
        );

        logger.log_video_frame(&stream_handle, img.clone(), 1);
        logger.log_video_frame(&stream_handle, img, 1);
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(2));

        logger.reset_video_stream(&stream_handle);
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(0));
    }

    #[test]
    fn test_video_stream_handle_drop() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let stream_handle = logger.init_video_stream("test_stream");
        let stream_key = stream_handle.key.clone();

        // Verify stream exists
        assert_eq!(logger.get_stream_frame_index(&stream_handle), Some(0));

        // Drop the handle
        drop(stream_handle);

        // Stream should be cleaned up - verify by checking the internal state
        let streams = logger.video_streams.lock().unwrap();
        assert!(!streams.contains_key(&stream_key));
    }

    #[test]
    fn test_multiple_video_streams() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let stream1 = logger.init_video_stream("stream1");
        let stream2 = logger.init_video_stream("stream2");

        assert_ne!(stream1.key, stream2.key);

        let img = rerun::Image::from_color_model_and_bytes(
            vec![0u8; 8 * 8 * 3],
            [8, 8],
            rerun::ColorModel::RGB,
            rerun::ChannelDatatype::U8,
        );

        logger.log_video_frame(&stream1, img.clone(), 1);
        logger.log_video_frame(&stream1, img.clone(), 1);
        logger.log_video_frame(&stream2, img, 1);

        assert_eq!(logger.get_stream_frame_index(&stream1), Some(2));
        assert_eq!(logger.get_stream_frame_index(&stream2), Some(1));
    }

    #[test]
    fn test_video_stream_persistence() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        // Initialize the same stream twice - should return handles to the same underlying stream
        let stream1 = logger.init_video_stream("persistent");
        let stream2 = logger.init_video_stream("persistent");

        assert_eq!(stream1.key, stream2.key);
        assert_eq!(logger.get_stream_frame_index(&stream1), Some(0));
    }

    #[test]
    fn test_record_to_path_creates_rrd_file() {
        let dir = tempdir().expect("temp dir");
        let record_path = dir.path().join("stream.rrd");

        let config = RerunConfig::default()
            .with_spawn_viewer(false)
            .with_record_to_path(record_path.clone());

        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");
        logger.log_scalar("test_metric", 1.0, 1);
        drop(logger);

        assert!(
            record_path.exists(),
            "RRD file should exist at {}",
            record_path.display()
        );
        assert!(
            record_path.metadata().unwrap().len() > 0,
            "RRD file should contain data"
        );
    }

    #[test]
    fn test_record_to_path_creates_parent_dirs() {
        let dir = tempdir().expect("temp dir");
        let nested = dir.path().join("nested/path/stream.rrd");

        let config = RerunConfig::default()
            .with_spawn_viewer(false)
            .with_record_to_path(nested.clone());

        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");
        logger.log_scalar("metric", 1.0, 1);
        drop(logger);

        assert!(
            nested.exists(),
            "RRD file should exist at {}",
            nested.display()
        );
    }

    #[test]
    fn test_video_frame_timelines_are_monotonic() {
        let config = RerunConfig::default().with_spawn_viewer(false);
        let logger = RerunLogger::new(config, "test-job").expect("Failed to create logger");

        let stream = logger.init_video_stream("monotonic");
        let img = rerun::Image::from_color_model_and_bytes(
            vec![0u8; 4],
            [1, 2],
            rerun::ColorModel::RGB,
            rerun::ChannelDatatype::U8,
        );

        logger.log_video_frame(&stream, img.clone(), 0);
        logger.log_video_frame(&stream, img.clone(), 0);
        logger.log_video_frame(&stream, img, 1);

        assert_eq!(logger.get_stream_frame_index(&stream), Some(3));
    }

    #[test]
    fn test_entity_path_resolution_with_explicit_path() {
        assert_eq!(
            resolve_entity_path("metrics/train/loss", "metrics"),
            "metrics/train/loss"
        );
    }

    #[test]
    fn test_entity_path_resolution_with_simple_name() {
        assert_eq!(
            resolve_entity_path("loss value", "metrics"),
            "metrics/loss_value"
        );
    }

    #[test]
    fn test_entity_path_resolution_on_empty_name() {
        assert_eq!(resolve_entity_path("", "metrics"), "metrics");
    }
}