profile-inspect 0.1.3

use std::collections::HashMap;
use std::io::Write;

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

use super::{Formatter, OutputError};

/// Collapsed stacks formatter for flamegraph tools
///
/// Output format (Brendan Gregg's folded format):
/// ```text
/// main;foo;bar 42
/// main;foo;baz 17
/// ```
///
/// Compatible with:
/// - Brendan Gregg's flamegraph.pl
/// - inferno (Rust flamegraph tool)
/// - speedscope (can import collapsed stacks)
pub struct CollapsedFormatter;

impl Formatter for CollapsedFormatter {
    fn write_cpu_analysis(
        &self,
        profile: &ProfileIR,
        _analysis: &CpuAnalysis,
        writer: &mut dyn Write,
    ) -> Result<(), OutputError> {
        // Aggregate stacks
        let mut stack_weights: HashMap<String, u64> = HashMap::new();

        for sample in &profile.samples {
            if let Some(stack) = profile.get_stack(sample.stack_id) {
                // Build collapsed stack string
                let stack_str: String = stack
                    .frames
                    .iter()
                    .filter_map(|fid| {
                        profile.get_frame(*fid).map(|f| {
                            // Clean up function name for collapsed format
                            let name = if f.name.is_empty() {
                                "(anonymous)".to_string()
                            } else {
                                // Replace semicolons which are used as delimiters
                                f.name.replace(';', ":")
                            };

                            // Optionally include file info
                            if let Some(ref file) = f.file {
                                // Extract just filename
                                let filename = file.rsplit('/').next().unwrap_or(file);
                                if let Some(line) = f.line {
                                    format!("{name} ({filename}:{line})")
                                } else {
                                    format!("{name} ({filename})")
                                }
                            } else {
                                name
                            }
                        })
                    })
                    .collect::<Vec<_>>()
                    .join(";");

                if !stack_str.is_empty() {
                    *stack_weights.entry(stack_str).or_default() += sample.weight;
                }
            }
        }

        // Sort by stack name for deterministic output
        let mut stacks: Vec<_> = stack_weights.into_iter().collect();
        stacks.sort_by(|a, b| a.0.cmp(&b.0));

        // Write collapsed stacks
        for (stack, weight) in stacks {
            writeln!(writer, "{stack} {weight}")?;
        }

        Ok(())
    }

    fn write_heap_analysis(
        &self,
        profile: &ProfileIR,
        _analysis: &HeapAnalysis,
        writer: &mut dyn Write,
    ) -> Result<(), OutputError> {
        // Same format but with allocation sizes instead of time
        let mut stack_weights: HashMap<String, u64> = HashMap::new();

        for sample in &profile.samples {
            if let Some(stack) = profile.get_stack(sample.stack_id) {
                let stack_str: String = stack
                    .frames
                    .iter()
                    .filter_map(|fid| {
                        profile.get_frame(*fid).map(|f| {
                            let name = if f.name.is_empty() {
                                "(anonymous)".to_string()
                            } else {
                                f.name.replace(';', ":")
                            };

                            if let Some(ref file) = f.file {
                                let filename = file.rsplit('/').next().unwrap_or(file);
                                if let Some(line) = f.line {
                                    format!("{name} ({filename}:{line})")
                                } else {
                                    format!("{name} ({filename})")
                                }
                            } else {
                                name
                            }
                        })
                    })
                    .collect::<Vec<_>>()
                    .join(";");

                if !stack_str.is_empty() {
                    *stack_weights.entry(stack_str).or_default() += sample.weight;
                }
            }
        }

        let mut stacks: Vec<_> = stack_weights.into_iter().collect();
        stacks.sort_by(|a, b| a.0.cmp(&b.0));

        for (stack, weight) in stacks {
            writeln!(writer, "{stack} {weight}")?;
        }

        Ok(())
    }
}