profile-inspect 0.1.3

use std::io::Write;

use indexmap::IndexMap;
use serde::Serialize;

use crate::analysis::{CpuAnalysis, HeapAnalysis};
use crate::ir::ProfileIR;

use super::{Formatter, OutputError};

/// Speedscope JSON formatter for visualization
pub struct SpeedscopeFormatter;

/// Speedscope file format
/// See: https://www.speedscope.app/file-format-spec.pdf
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct SpeedscopeFile {
    #[serde(rename = "$schema")]
    schema: &'static str,
    shared: SpeedscopeShared,
    profiles: Vec<SpeedscopeProfile>,
    name: Option<String>,
    active_profile_index: Option<usize>,
    exporter: Option<String>,
}

#[derive(Serialize)]
struct SpeedscopeShared {
    frames: Vec<SpeedscopeFrame>,
}

#[derive(Serialize)]
struct SpeedscopeFrame {
    name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    file: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    line: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    col: Option<u32>,
}

#[derive(Serialize)]
#[serde(tag = "type")]
enum SpeedscopeProfile {
    #[serde(rename = "sampled")]
    Sampled {
        name: String,
        unit: &'static str,
        #[serde(rename = "startValue")]
        start_value: u64,
        #[serde(rename = "endValue")]
        end_value: u64,
        samples: Vec<Vec<usize>>,
        weights: Vec<u64>,
    },
}

impl Formatter for SpeedscopeFormatter {
    fn write_cpu_analysis(
        &self,
        profile: &ProfileIR,
        analysis: &CpuAnalysis,
        writer: &mut dyn Write,
    ) -> Result<(), OutputError> {
        // Build frame index (map from FrameId to speedscope index)
        let mut frame_to_index: IndexMap<crate::ir::FrameId, usize> = IndexMap::new();
        let mut speedscope_frames: Vec<SpeedscopeFrame> = Vec::new();

        for frame in &profile.frames {
            let idx = speedscope_frames.len();
            frame_to_index.insert(frame.id, idx);

            speedscope_frames.push(SpeedscopeFrame {
                name: if frame.name.is_empty() {
                    "(anonymous)".to_string()
                } else {
                    frame.name.clone()
                },
                file: frame.clean_file(),
                line: frame.line,
                col: frame.col,
            });
        }

        // Build samples and weights
        let mut samples: Vec<Vec<usize>> = Vec::with_capacity(profile.samples.len());
        let mut weights: Vec<u64> = Vec::with_capacity(profile.samples.len());

        for sample in &profile.samples {
            if let Some(stack) = profile.get_stack(sample.stack_id) {
                let sample_indices: Vec<usize> = stack
                    .frames
                    .iter()
                    .filter_map(|fid| frame_to_index.get(fid).copied())
                    .collect();

                if !sample_indices.is_empty() {
                    samples.push(sample_indices);
                    weights.push(sample.weight);
                }
            }
        }

        let file = SpeedscopeFile {
            schema: "https://www.speedscope.app/file-format-schema.json",
            shared: SpeedscopeShared {
                frames: speedscope_frames,
            },
            profiles: vec![SpeedscopeProfile::Sampled {
                name: profile
                    .source_file
                    .clone()
                    .unwrap_or_else(|| "CPU Profile".to_string()),
                unit: "microseconds",
                start_value: 0,
                end_value: analysis.total_time,
                samples,
                weights,
            }],
            name: profile.source_file.clone(),
            active_profile_index: Some(0),
            exporter: Some("profile-inspect".to_string()),
        };

        serde_json::to_writer(writer, &file)?;
        Ok(())
    }

    fn write_heap_analysis(
        &self,
        profile: &ProfileIR,
        analysis: &HeapAnalysis,
        writer: &mut dyn Write,
    ) -> Result<(), OutputError> {
        // Build frame index
        let mut frame_to_index: IndexMap<crate::ir::FrameId, usize> = IndexMap::new();
        let mut speedscope_frames: Vec<SpeedscopeFrame> = Vec::new();

        for frame in &profile.frames {
            let idx = speedscope_frames.len();
            frame_to_index.insert(frame.id, idx);

            speedscope_frames.push(SpeedscopeFrame {
                name: if frame.name.is_empty() {
                    "(anonymous)".to_string()
                } else {
                    frame.name.clone()
                },
                file: frame.clean_file(),
                line: frame.line,
                col: frame.col,
            });
        }

        // Build samples and weights (bytes instead of time)
        let mut samples: Vec<Vec<usize>> = Vec::with_capacity(profile.samples.len());
        let mut weights: Vec<u64> = Vec::with_capacity(profile.samples.len());

        for sample in &profile.samples {
            if let Some(stack) = profile.get_stack(sample.stack_id) {
                let sample_indices: Vec<usize> = stack
                    .frames
                    .iter()
                    .filter_map(|fid| frame_to_index.get(fid).copied())
                    .collect();

                if !sample_indices.is_empty() {
                    samples.push(sample_indices);
                    weights.push(sample.weight);
                }
            }
        }

        let file = SpeedscopeFile {
            schema: "https://www.speedscope.app/file-format-schema.json",
            shared: SpeedscopeShared {
                frames: speedscope_frames,
            },
            profiles: vec![SpeedscopeProfile::Sampled {
                name: profile
                    .source_file
                    .clone()
                    .unwrap_or_else(|| "Heap Profile".to_string()),
                unit: "bytes",
                start_value: 0,
                end_value: analysis.total_size,
                samples,
                weights,
            }],
            name: profile.source_file.clone(),
            active_profile_index: Some(0),
            exporter: Some("profile-inspect".to_string()),
        };

        serde_json::to_writer(writer, &file)?;
        Ok(())
    }
}