ssb 0.1.1

use std::collections::HashMap;
use std::io::IsTerminal as _;
use std::path::Path;

use fastrace::collector::SpanId;
use serde::{Deserialize, Serialize};

use crate::bench::BENCH_ROOT;
use crate::collector::RawSpan;
use crate::stats::{SpanStats, fmt_ns};

// ── ANSI colours ─────────────────────────────────────────────────────────────

const GREEN: &str = "\x1b[32m";
const RED: &str = "\x1b[31m";
const RESET: &str = "\x1b[0m";

fn use_color() -> bool {
    std::io::stdout().is_terminal()
}

// ── public types ─────────────────────────────────────────────────────────────

/// A complete benchmark report: per-span statistics across all iterations.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchReport {
    pub name: String,
    pub iterations: usize,
    /// Unix timestamp (seconds) of when this report was produced.
    pub timestamp_unix_secs: u64,
    pub spans: HashMap<String, SpanStats>,
}

impl BenchReport {
    pub(crate) fn from_iters(name: String, all_iters: Vec<Vec<RawSpan>>) -> Self {
        let n = all_iters.len();
        let mut per_span: HashMap<String, Vec<u64>> = HashMap::new();
        // First-seen parent name for each display span name.
        let mut parent_names: HashMap<String, Option<String>> = HashMap::new();

        for iter in all_iters {
            let id_to_name: HashMap<SpanId, String> = iter
                .iter()
                .map(|raw| (raw.span_id, display_name(&raw.name, &name)))
                .collect();

            let mut totals: HashMap<String, u64> = HashMap::new();
            for raw in &iter {
                *totals.entry(display_name(&raw.name, &name)).or_default() += raw.duration_ns;
            }
            for (span_name, dur) in totals {
                per_span.entry(span_name).or_default().push(dur);
            }

            for raw in &iter {
                let display = display_name(&raw.name, &name);
                parent_names
                    .entry(display)
                    .or_insert_with(|| id_to_name.get(&raw.parent_id).cloned());
            }
        }

        for parent in parent_names.values_mut() {
            if let Some(ref p) = *parent
                && !per_span.contains_key(p.as_str())
            {
                *parent = None;
            }
        }

        let spans: HashMap<String, SpanStats> = per_span
            .into_iter()
            .map(|(k, v)| {
                let parent = parent_names.remove(&k).flatten();
                let stats = SpanStats::compute(k.clone(), v, parent);
                (k, stats)
            })
            .collect();

        let timestamp_unix_secs = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        BenchReport {
            name,
            iterations: n,
            timestamp_unix_secs,
            spans,
        }
    }

    /// Print a formatted tree of span statistics to stdout.
    pub fn print(&self) {
        let children = build_children(&self.spans);
        let order = dfs_order(&children, &self.name);
        let (name_w, col_w, pct_w) = stats_widths(&order);
        let sep = stats_sep(name_w, col_w, pct_w);

        println!();
        println!("Benchmark: {} ({} iterations)", self.name, self.iterations);
        print_stats_header(name_w, col_w, pct_w, &sep);
        print_stats_rows(&order, &self.spans, name_w, col_w, pct_w);
        println!("{sep}");
        println!();
    }

    /// Serialize the report to a JSON file for later comparison.
    pub fn save(&self, path: impl AsRef<Path>) -> std::io::Result<()> {
        let json = serde_json::to_string_pretty(self).map_err(std::io::Error::other)?;
        std::fs::write(path, json)
    }

    /// Load a previously saved report from a JSON file.
    pub fn load(path: impl AsRef<Path>) -> std::io::Result<Self> {
        let text = std::fs::read_to_string(path)?;
        serde_json::from_str(&text).map_err(std::io::Error::other)
    }

    /// Build a [`Comparison`] against a baseline report.
    ///
    /// `self` is the *current* (new) result; `baseline` is the reference.
    pub fn compare<'a>(&'a self, baseline: &'a BenchReport) -> Comparison<'a> {
        Comparison {
            baseline,
            current: self,
        }
    }
}

/// Side-by-side comparison of a current report against a baseline.
pub struct Comparison<'a> {
    pub baseline: &'a BenchReport,
    pub current: &'a BenchReport,
}

impl Comparison<'_> {
    /// Print a formatted comparison table to stdout.
    pub fn print(&self) {
        let order = unified_order(self.baseline, self.current);
        let (name_w, col_w, diff_w) = cmp_widths(&order);
        let sep = cmp_sep(name_w, col_w, diff_w);

        println!();
        println!("Comparison: {}", self.baseline.name);
        println!("  baseline: {} iterations", self.baseline.iterations);
        println!("  current:  {} iterations", self.current.iterations);
        print_cmp_header(name_w, col_w, diff_w, &sep);
        print_cmp_rows(&order, self.baseline, self.current, name_w, col_w, diff_w);
        println!("{sep}");
        println!();
    }
}

// ── group printing (called from bench.rs) ────────────────────────────────────

/// Print a combined table for all benchmarks in a group.
///
/// `entries` is a slice of `(current, baseline_or_none)` pairs in group order.
pub(crate) fn print_group(entries: &[(&BenchReport, Option<&BenchReport>)], group_name: &str) {
    if entries.is_empty() {
        return;
    }

    let has_all = entries.iter().all(|(_, b)| b.is_some());
    let has_none = entries.iter().all(|(_, b)| b.is_none());

    if !has_all && !has_none {
        // Mixed — fall back to printing individually.
        for (report, baseline) in entries {
            match baseline {
                Some(b) => report.compare(b).print(),
                None => report.print(),
            }
        }
        return;
    }

    if has_none {
        // ── Stats table ──────────────────────────────────────────────────────
        let orders: Vec<Vec<(&str, usize)>> = entries
            .iter()
            .map(|(r, _)| dfs_order(&build_children(&r.spans), &r.name))
            .collect();

        // Compute widths across all entries.
        let all_order: Vec<(&&str, &usize)> = orders
            .iter()
            .flat_map(|o| o.iter().map(|(n, d)| (n, d)))
            .collect();
        let (name_w, col_w, pct_w) = stats_widths_raw(&all_order, &orders);
        let sep = stats_sep(name_w, col_w, pct_w);
        let thin = thin_sep(sep.chars().count());

        println!();
        println!(
            "Benchmark Group: {}  ({} iterations)",
            group_name, entries[0].0.iterations,
        );
        print_stats_header(name_w, col_w, pct_w, &sep);

        for (i, ((report, _), order)) in entries.iter().zip(orders.iter()).enumerate() {
            if i > 0 {
                println!("{thin}");
            }
            print_stats_rows(order, &report.spans, name_w, col_w, pct_w);
        }
        println!("{sep}");
        println!();
    } else {
        // ── Comparison table ─────────────────────────────────────────────────
        let orders: Vec<Vec<(&str, usize)>> = entries
            .iter()
            .map(|(c, b)| unified_order(b.unwrap(), c))
            .collect();

        let all_order: Vec<(&&str, &usize)> = orders
            .iter()
            .flat_map(|o| o.iter().map(|(n, d)| (n, d)))
            .collect();
        let max_name = all_order
            .iter()
            .map(|(name, _)| name.len())
            .max()
            .unwrap_or(4);
        // For cross-entry prefix computation we compute max depth separately.
        let max_prefix = orders
            .iter()
            .flat_map(|o| {
                o.iter()
                    .enumerate()
                    .map(|(idx, _)| tree_prefix(o, idx).chars().count())
            })
            .max()
            .unwrap_or(0);
        let name_w = (max_prefix + max_name).max(4) + 2;
        let col_w: usize = 11;
        let diff_w: usize = 20;
        let sep = cmp_sep(name_w, col_w, diff_w);
        let thin = thin_sep(sep.chars().count());

        println!();
        println!("Comparison Group: {}", group_name);
        println!(
            "  baseline: {} iterations",
            entries[0].1.unwrap().iterations
        );
        println!("  current:  {} iterations", entries[0].0.iterations);
        print_cmp_header(name_w, col_w, diff_w, &sep);

        for (i, ((current, baseline_opt), order)) in entries.iter().zip(orders.iter()).enumerate() {
            if i > 0 {
                println!("{thin}");
            }
            print_cmp_rows(order, baseline_opt.unwrap(), current, name_w, col_w, diff_w);
        }
        println!("{sep}");
        println!();
    }
}

// ── separator / width helpers ─────────────────────────────────────────────────

fn stats_sep(name_w: usize, col_w: usize, pct_w: usize) -> String {
    "─".repeat(name_w + col_w * 4 + pct_w + 12)
}

fn cmp_sep(name_w: usize, col_w: usize, diff_w: usize) -> String {
    "─".repeat(name_w + col_w * 2 + diff_w + 12)
}

/// Width of the name column and fixed column widths for stats tables.
fn stats_widths(order: &[(&str, usize)]) -> (usize, usize, usize) {
    let col_w: usize = 11;
    let pct_w: usize = 9;
    let name_w = order
        .iter()
        .enumerate()
        .map(|(idx, (name, _))| tree_prefix(order, idx).chars().count() + name.len())
        .max()
        .unwrap_or(4)
        .max(4)
        + 2;
    (name_w, col_w, pct_w)
}

/// Width calculation when we have multiple orders (group case).
/// `all_order` is a flat view used only for name length; `orders` is used for
/// prefix length per entry.
fn stats_widths_raw(
    _all_order: &[(&&str, &usize)],
    orders: &[Vec<(&str, usize)>],
) -> (usize, usize, usize) {
    let col_w: usize = 11;
    let pct_w: usize = 9;
    let name_w = orders
        .iter()
        .flat_map(|o| {
            o.iter()
                .enumerate()
                .map(|(idx, (name, _))| tree_prefix(o, idx).chars().count() + name.len())
        })
        .max()
        .unwrap_or(4)
        .max(4)
        + 2;
    (name_w, col_w, pct_w)
}

/// Width of the name column and fixed column widths for comparison tables.
fn cmp_widths(order: &[(&str, usize)]) -> (usize, usize, usize) {
    let col_w: usize = 11;
    let diff_w: usize = 20;
    let name_w = order
        .iter()
        .enumerate()
        .map(|(idx, (name, _))| tree_prefix(order, idx).chars().count() + name.len())
        .max()
        .unwrap_or(4)
        .max(4)
        + 2;
    (name_w, col_w, diff_w)
}

fn thin_sep(width: usize) -> String {
    (0..width)
        .map(|i| if i % 2 == 0 { '─' } else { ' ' })
        .collect()
}

// ── table printing ────────────────────────────────────────────────────────────

fn print_stats_header(name_w: usize, col_w: usize, pct_w: usize, sep: &str) {
    println!("{sep}");
    println!(
        " {:<nw$}  {:>cw$}  {:>cw$}  {:>cw$}  {:>cw$}  {:>pw$}",
        "span",
        "min",
        "median",
        "p95",
        "max",
        "% parent",
        nw = name_w,
        cw = col_w,
        pw = pct_w,
    );
    println!("{sep}");
}

fn print_stats_rows(
    order: &[(&str, usize)],
    spans: &HashMap<String, SpanStats>,
    name_w: usize,
    col_w: usize,
    pct_w: usize,
) {
    for (idx, (name, _)) in order.iter().enumerate() {
        let s = &spans[*name];
        let prefix = tree_prefix(order, idx);
        let indented = format!("{prefix}{name}");
        let pct_str = pct_of_parent(s, spans);
        println!(
            " {:<nw$}  {:>cw$}  {:>cw$}  {:>cw$}  {:>cw$}  {:>pw$}",
            indented,
            fmt_ns(s.min_ns as f64),
            fmt_ns(s.median_ns),
            fmt_ns(s.p95_ns),
            fmt_ns(s.max_ns as f64),
            pct_str,
            nw = name_w,
            cw = col_w,
            pw = pct_w,
        );
    }
}

fn print_cmp_header(name_w: usize, col_w: usize, diff_w: usize, sep: &str) {
    println!("{sep}");
    println!(
        " {:<nw$}  {:>cw$}  {:>cw$}  {:<dw$}",
        "span",
        "baseline",
        "current",
        "change (median)",
        nw = name_w,
        cw = col_w,
        dw = diff_w,
    );
    println!("{sep}");
}

fn print_cmp_rows(
    order: &[(&str, usize)],
    baseline: &BenchReport,
    current: &BenchReport,
    name_w: usize,
    col_w: usize,
    diff_w: usize,
) {
    let colored = use_color();

    for (idx, (name, _)) in order.iter().enumerate() {
        let bs = baseline.spans.get(*name);
        let cs = current.spans.get(*name);

        let prefix = tree_prefix(order, idx);
        let indented = format!("{prefix}{name}");
        let base_str = bs
            .map(|s| fmt_ns(s.median_ns))
            .unwrap_or_else(|| "n/a".into());
        let curr_str = cs
            .map(|s| fmt_ns(s.median_ns))
            .unwrap_or_else(|| "n/a".into());

        let (diff_plain, color) = match (bs, cs) {
            (Some(bs), Some(cs)) => {
                let pct = (cs.median_ns - bs.median_ns) / bs.median_ns * 100.0;
                if pct <= -1.0 {
                    (format!("{:+.1}% ↓  faster", pct), GREEN)
                } else if pct >= 1.0 {
                    (format!("{:+.1}% ↑  slower", pct), RED)
                } else {
                    (format!("{:+.2}%    unchanged", pct), "")
                }
            }
            (None, Some(_)) => ("new".to_string(), GREEN),
            (Some(_), None) => ("removed".to_string(), RED),
            (None, None) => unreachable!(),
        };

        // Build diff cell with colour; pad manually so ANSI codes don't throw
        // off Rust's format! padding (which counts bytes, not display width).
        let plain_len = diff_plain.chars().count();
        let padding = " ".repeat(diff_w.saturating_sub(plain_len));
        let diff_cell = if colored && !color.is_empty() {
            format!("{color}{diff_plain}{RESET}{padding}")
        } else {
            format!("{diff_plain}{padding}")
        };

        println!(
            " {:<nw$}  {:>cw$}  {:>cw$}  {}",
            indented,
            base_str,
            curr_str,
            diff_cell,
            nw = name_w,
            cw = col_w,
        );
    }
}

// ── tree helpers ──────────────────────────────────────────────────────────────

/// Returns true if there is another node at the same depth as `order[pos]`
/// after position `pos`, before any shallower node.
fn has_more_sibling(order: &[(&str, usize)], pos: usize) -> bool {
    let target = order[pos].1;
    for &(_, d) in &order[pos + 1..] {
        if d < target {
            return false;
        }
        if d == target {
            return true;
        }
    }
    false
}

/// Build the tree-drawing prefix string for the node at `idx` in DFS `order`.
///
/// Each depth level uses 4 display characters:
///   `│   ` or `    ` for ancestor continuation,
///   `├── ` or `└── ` for the node's own connector.
fn tree_prefix(order: &[(&str, usize)], idx: usize) -> String {
    let depth = order[idx].1;
    if depth == 0 {
        return String::new();
    }
    let mut prefix = String::new();
    for level in 0..depth {
        let ref_pos = if level == depth - 1 {
            idx
        } else {
            // Nearest ancestor at depth `level + 1`.
            order[..idx]
                .iter()
                .rposition(|&(_, d)| d == level + 1)
                .unwrap_or(idx)
        };
        let has_more = has_more_sibling(order, ref_pos);
        if level == depth - 1 {
            prefix.push_str(if has_more { "├── " } else { "└── " });
        } else {
            prefix.push_str(if has_more { "│   " } else { "    " });
        }
    }
    prefix
}

// ── span tree / DFS helpers ───────────────────────────────────────────────────

fn display_name(raw_name: &str, bench_name: &str) -> String {
    if raw_name == BENCH_ROOT {
        bench_name.to_owned()
    } else {
        raw_name.to_owned()
    }
}

fn build_children(spans: &HashMap<String, SpanStats>) -> HashMap<Option<&'_ str>, Vec<&'_ str>> {
    let bench_name: &str = spans
        .values()
        .find(|s| s.parent.is_none())
        .map(|s| s.name.as_str())
        .unwrap_or("");

    let mut children: HashMap<Option<&str>, Vec<&str>> = HashMap::new();
    for (name, stats) in spans {
        children
            .entry(stats.parent.as_deref())
            .or_default()
            .push(name.as_str());
    }
    for (_, kids) in children.iter_mut() {
        kids.sort_by(|a, b| sort_key(a, b, bench_name));
    }
    children
}

fn sort_key(a: &&str, b: &&str, bench_name: &str) -> std::cmp::Ordering {
    match (a == &bench_name, b == &bench_name) {
        (true, false) => std::cmp::Ordering::Less,
        (false, true) => std::cmp::Ordering::Greater,
        _ => a.cmp(b),
    }
}

fn dfs_order<'a>(
    children: &HashMap<Option<&'a str>, Vec<&'a str>>,
    bench_name: &'a str,
) -> Vec<(&'a str, usize)> {
    dfs_from_children(children, bench_name)
}

fn dfs_from_children<'a>(
    children: &HashMap<Option<&'a str>, Vec<&'a str>>,
    _bench_name: &'a str,
) -> Vec<(&'a str, usize)> {
    let mut order = Vec::new();
    let mut stack: Vec<(&str, usize)> = Vec::new();

    if let Some(roots) = children.get(&None) {
        for &r in roots.iter().rev() {
            stack.push((r, 0));
        }
    }

    let mut visited = std::collections::HashSet::new();
    while let Some((name, depth)) = stack.pop() {
        if !visited.insert(name) {
            continue;
        }
        order.push((name, depth));
        if let Some(kids) = children.get(&Some(name)) {
            for &kid in kids.iter().rev() {
                if !visited.contains(kid) {
                    stack.push((kid, depth + 1));
                }
            }
        }
    }

    // Catch unreachable spans.
    let all_names: std::collections::HashSet<&str> = children.values().flatten().copied().collect();
    for name in all_names {
        if !visited.contains(name) {
            order.push((name, 0));
        }
    }

    order
}

/// Build a unified DFS order merging spans from both baseline and current.
fn unified_order<'a>(baseline: &'a BenchReport, current: &'a BenchReport) -> Vec<(&'a str, usize)> {
    let bench_name = baseline.name.as_str();
    let mut parent_of: HashMap<&str, Option<&str>> = HashMap::new();
    for (name, stats) in &current.spans {
        parent_of.insert(name.as_str(), stats.parent.as_deref());
    }
    for (name, stats) in &baseline.spans {
        parent_of
            .entry(name.as_str())
            .or_insert(stats.parent.as_deref());
    }

    let mut children: HashMap<Option<&str>, Vec<&str>> = HashMap::new();
    for (&name, &parent) in &parent_of {
        children.entry(parent).or_default().push(name);
    }
    for kids in children.values_mut() {
        kids.sort_by(|a, b| sort_key(a, b, bench_name));
    }
    dfs_from_children(&children, bench_name)
}

fn pct_of_parent(s: &SpanStats, spans: &HashMap<String, SpanStats>) -> String {
    let Some(ref parent_name) = s.parent else {
        return "—".to_string();
    };
    let Some(parent) = spans.get(parent_name) else {
        return "—".to_string();
    };
    if parent.mean_ns <= 0.0 {
        return "—".to_string();
    }
    let pct = (s.mean_ns / parent.mean_ns * 100.0).round() as u64;
    format!("{pct}%")
}