jtool_notebook/

timing.rs

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Execution timing metadata from notebook cells
///
/// JupyterLab 1.2+ can record timing information in cell metadata under the "execution" key.
/// This captures timestamps at various stages of cell execution.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ExecutionTiming {
    /// When kernel became busy (started processing request)
    #[serde(rename = "iopub.status.busy", skip_serializing_if = "Option::is_none")]
    pub status_busy: Option<DateTime<Utc>>,

    /// When execute_input message was sent
    #[serde(
        rename = "iopub.execute_input",
        skip_serializing_if = "Option::is_none"
    )]
    pub execute_input: Option<DateTime<Utc>>,

    /// When execution started (from shell.execute_reply metadata)
    #[serde(
        rename = "shell.execute_reply.started",
        skip_serializing_if = "Option::is_none"
    )]
    pub reply_started: Option<DateTime<Utc>>,

    /// When execution finished (shell.execute_reply received)
    #[serde(
        rename = "shell.execute_reply",
        skip_serializing_if = "Option::is_none"
    )]
    pub reply: Option<DateTime<Utc>>,

    /// When kernel became idle (all outputs published)
    #[serde(rename = "iopub.status.idle", skip_serializing_if = "Option::is_none")]
    pub status_idle: Option<DateTime<Utc>>,
}

impl ExecutionTiming {
    /// Calculate total duration (busy → idle)
    ///
    /// This is the wall-clock time from when the kernel started processing
    /// the request until all outputs were published.
    pub fn total_duration(&self) -> Option<Duration> {
        match (self.status_busy, self.status_idle) {
            (Some(busy), Some(idle)) => {
                let duration = idle.signed_duration_since(busy);
                duration.to_std().ok()
            }
            _ => None,
        }
    }

    /// Calculate pure execution duration (reply_started → reply)
    ///
    /// This is the actual execution time, excluding output publishing overhead.
    pub fn execution_duration(&self) -> Option<Duration> {
        match (self.reply_started, self.reply) {
            (Some(started), Some(finished)) => {
                let duration = finished.signed_duration_since(started);
                duration.to_std().ok()
            }
            _ => None,
        }
    }

    /// Calculate output overhead (reply → idle)
    ///
    /// This is the time spent publishing outputs after execution finished.
    pub fn output_overhead(&self) -> Option<Duration> {
        match (self.reply, self.status_idle) {
            (Some(reply), Some(idle)) => {
                let duration = idle.signed_duration_since(reply);
                duration.to_std().ok()
            }
            _ => None,
        }
    }

    /// Create timing from start and end timestamps
    ///
    /// This is a simplified version for when we only have start/end times
    /// (e.g., when capturing timing during execution without all intermediate events)
    pub fn from_duration(start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
        Self {
            status_busy: Some(start),
            execute_input: Some(start),
            reply_started: Some(start),
            reply: Some(end),
            status_idle: Some(end),
        }
    }

    /// Check if this timing has valid data
    pub fn has_timing(&self) -> bool {
        self.status_busy.is_some() && self.status_idle.is_some()
    }
}

/// Cell timing information with index and execution count
#[derive(Debug, Clone, Serialize)]
pub struct CellTiming {
    /// Cell index in the notebook
    pub cell_index: usize,
    /// Execution count (if the cell has been executed)
    pub execution_count: Option<i32>,
    /// Timing metadata
    pub timing: ExecutionTiming,
    /// Calculated duration (total_duration from timing)
    pub duration: Option<Duration>,
}

impl CellTiming {
    /// Create a new CellTiming
    pub fn new(cell_index: usize, execution_count: Option<i32>, timing: ExecutionTiming) -> Self {
        let duration = timing.total_duration();
        Self {
            cell_index,
            execution_count,
            timing,
            duration,
        }
    }

    /// Get duration as seconds (f64)
    pub fn duration_seconds(&self) -> Option<f64> {
        self.duration.map(|d| d.as_secs_f64())
    }
}

/// Summary statistics for notebook timing
#[derive(Debug, Clone, Serialize)]
pub struct TimingSummary {
    /// Total execution time across all cells
    pub total_duration: Duration,
    /// Number of cells that have been executed
    pub cells_executed: usize,
    /// Average execution time per cell
    pub average_duration: Duration,
    /// Maximum execution time
    pub max_duration: Duration,
    /// Cell index with maximum execution time
    pub max_duration_cell: usize,
    /// Minimum execution time
    pub min_duration: Duration,
    /// Cell index with minimum execution time
    pub min_duration_cell: usize,
}

impl TimingSummary {
    /// Calculate summary from a list of cell timings
    pub fn from_cells(cells: &[CellTiming]) -> Option<Self> {
        if cells.is_empty() {
            return None;
        }

        let cells_with_duration: Vec<_> = cells.iter().filter(|c| c.duration.is_some()).collect();

        if cells_with_duration.is_empty() {
            return None;
        }

        let total_duration: Duration = cells_with_duration.iter().filter_map(|c| c.duration).sum();

        let average_duration = total_duration
            .checked_div(cells_with_duration.len() as u32)
            .unwrap_or(Duration::ZERO);

        // Find max and min durations
        // Safe: cells_with_duration is non-empty (checked above) and all have Some(duration)
        let first_cell = cells_with_duration.first()?;

        let (max_cell, min_cell) =
            cells_with_duration
                .iter()
                .fold((first_cell, first_cell), |(max, min), cell| {
                    let new_max = if cell.duration > max.duration {
                        cell
                    } else {
                        max
                    };
                    let new_min = if cell.duration < min.duration {
                        cell
                    } else {
                        min
                    };
                    (new_max, new_min)
                });

        // Safe: we know these cells have duration because of the filter above
        let max_duration = max_cell.duration?;
        let max_duration_cell = max_cell.cell_index;

        let min_duration = min_cell.duration?;
        let min_duration_cell = min_cell.cell_index;

        Some(Self {
            total_duration,
            cells_executed: cells_with_duration.len(),
            average_duration,
            max_duration,
            max_duration_cell,
            min_duration,
            min_duration_cell,
        })
    }
}

/// Format for displaying durations
#[derive(Debug, Clone)]
pub enum TimingFormat {
    /// Display as seconds with decimals (e.g., "1.234s")
    Seconds,
    /// Display as milliseconds (e.g., "1234ms")
    Milliseconds,
    /// Display in human-readable format (e.g., "2m 5s" or "1s 234ms")
    Human,
}

/// Format a duration for display
pub fn format_duration(duration: Duration, format: &TimingFormat) -> String {
    match format {
        TimingFormat::Seconds => {
            format!("{:.2}s", duration.as_secs_f64())
        }
        TimingFormat::Milliseconds => {
            format!("{}ms", duration.as_millis())
        }
        TimingFormat::Human => {
            let secs = duration.as_secs();
            let millis = duration.subsec_millis();

            if secs >= 60 {
                let mins = secs / 60;
                let remaining_secs = secs % 60;
                if millis > 0 {
                    format!("{mins}m {remaining_secs}s {millis}ms")
                } else {
                    format!("{mins}m {remaining_secs}s")
                }
            } else if secs > 0 {
                if millis > 0 {
                    format!("{secs}s {millis}ms")
                } else {
                    format!("{secs}s")
                }
            } else {
                format!("{millis}ms")
            }
        }
    }
}

/// Parse a duration string (e.g., "5s", "500ms", "2m")
pub fn parse_duration(s: &str) -> Result<Duration, String> {
    let s = s.trim();

    if let Some(stripped) = s.strip_suffix("ms") {
        let num = stripped
            .trim()
            .parse::<u64>()
            .map_err(|_| format!("Invalid duration: {s}"))?;
        Ok(Duration::from_millis(num))
    } else if let Some(stripped) = s.strip_suffix('s') {
        let num = stripped
            .trim()
            .parse::<f64>()
            .map_err(|_| format!("Invalid duration: {s}"))?;
        Ok(Duration::from_secs_f64(num))
    } else if let Some(stripped) = s.strip_suffix('m') {
        let num = stripped
            .trim()
            .parse::<u64>()
            .map_err(|_| format!("Invalid duration: {s}"))?;
        Ok(Duration::from_secs(num * 60))
    } else {
        Err(format!(
            "Invalid duration format: {s}. Use '5s', '500ms', or '2m'"
        ))
    }
}

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

    #[test]
    fn test_timing_total_duration() {
        let timing = ExecutionTiming {
            status_busy: "2020-01-01T00:00:00Z".parse().ok(),
            status_idle: "2020-01-01T00:00:05Z".parse().ok(),
            ..Default::default()
        };

        assert_eq!(timing.total_duration().map(|d| d.as_secs()), Some(5));
    }

    #[test]
    fn test_timing_execution_duration() {
        let timing = ExecutionTiming {
            reply_started: "2020-01-01T00:00:00Z".parse().ok(),
            reply: "2020-01-01T00:00:03Z".parse().ok(),
            ..Default::default()
        };

        assert_eq!(timing.execution_duration().map(|d| d.as_secs()), Some(3));
    }

    #[test]
    fn test_format_duration_seconds() {
        let duration = Duration::from_millis(1234);
        assert_eq!(format_duration(duration, &TimingFormat::Seconds), "1.23s");
    }

    #[test]
    fn test_format_duration_milliseconds() {
        let duration = Duration::from_millis(1234);
        assert_eq!(
            format_duration(duration, &TimingFormat::Milliseconds),
            "1234ms"
        );
    }

    #[test]
    fn test_format_duration_human() {
        let duration = Duration::from_secs(125);
        assert_eq!(format_duration(duration, &TimingFormat::Human), "2m 5s");

        let duration = Duration::from_millis(1500);
        assert_eq!(format_duration(duration, &TimingFormat::Human), "1s 500ms");

        let duration = Duration::from_millis(500);
        assert_eq!(format_duration(duration, &TimingFormat::Human), "500ms");
    }

    #[test]
    fn test_parse_duration() {
        assert_eq!(parse_duration("5s"), Ok(Duration::from_secs(5)));
        assert_eq!(parse_duration("500ms"), Ok(Duration::from_millis(500)));
        assert_eq!(parse_duration("2m"), Ok(Duration::from_secs(120)));
        assert_eq!(parse_duration("1.5s"), Ok(Duration::from_millis(1500)));

        assert!(parse_duration("invalid").is_err());
        assert!(parse_duration("5x").is_err());
    }

    #[test]
    fn test_timing_summary() {
        let cells = vec![
            CellTiming::new(
                0,
                Some(1),
                ExecutionTiming {
                    status_busy: "2020-01-01T00:00:00Z".parse().ok(),
                    status_idle: "2020-01-01T00:00:01Z".parse().ok(),
                    ..Default::default()
                },
            ),
            CellTiming::new(
                1,
                Some(2),
                ExecutionTiming {
                    status_busy: "2020-01-01T00:00:00Z".parse().ok(),
                    status_idle: "2020-01-01T00:00:05Z".parse().ok(),
                    ..Default::default()
                },
            ),
        ];

        let summary = TimingSummary::from_cells(&cells);
        assert!(summary.is_some(), "Failed to create summary: {summary:?}");
        if let Some(summary) = summary {
            assert_eq!(summary.total_duration.as_secs(), 6);
            assert_eq!(summary.cells_executed, 2);
            assert_eq!(summary.average_duration.as_secs(), 3);
            assert_eq!(summary.max_duration.as_secs(), 5);
            assert_eq!(summary.max_duration_cell, 1);
        }
    }
}