ssb 0.1.1

use std::path::PathBuf;

use fastrace::{Span, collector::SpanContext};

use crate::report::print_group;
use crate::{collector, report::BenchReport};

/// The fixed name used for the synthetic root span bencher creates per iteration.
/// Renamed to the benchmark name in the final report.
pub(crate) const BENCH_ROOT: &str = "__bencher_root__";

/// Builder for a single benchmark.
///
/// By default, [`Bench::run`] automatically:
/// 1. Loads `target/bencher/<name>/baseline.json` if it exists.
/// 2. Prints a comparison against that baseline (or just the stats on first run).
/// 3. Saves the new result as the next baseline.
///
/// Call [`Bench::no_auto_save`] to disable this and get the raw [`BenchReport`]
/// back without any I/O or printing side-effects.
///
/// # Thread safety
///
/// Bencher registers a global [`fastrace`] reporter on first use. Running
/// multiple `Bench::run` calls concurrently from different threads will mix
/// their span data. Run benchmarks sequentially (or use
/// `cargo test -- --test-threads=1`).
pub struct Bench {
    name: Option<String>,
    group: Option<String>,
    pub(crate) iterations: usize,
    pub(crate) min_run_seconds: usize,
    pub(crate) warmup_seconds: usize,
    pub(crate) auto_save: bool,
}

impl Bench {
    /// Create a new benchmark.
    ///
    /// Results are stored under `target/bencher/<name>/baseline.json` relative
    /// to the current working directory (set by cargo to the package/workspace
    /// root).
    pub fn new(name: impl Into<String>) -> Self {
        Bench {
            name: Some(name.into()),
            group: None,
            iterations: 1000,
            min_run_seconds: 3,
            warmup_seconds: 2,
            auto_save: true,
        }
    }

    /// Assign a group, organised into a subdirectory and printed as a single
    /// combined table. Returns a [`BenchGroup`] that collects all benchmarks
    /// and prints them together on drop.
    pub fn group(&mut self, group: impl Into<String>) -> BenchGroup<'_> {
        BenchGroup {
            bench: self,
            group_name: group.into(),
            entries: Vec::new(),
        }
    }

    /// Number of timed iterations (default: 1000).
    /// The measurement will run for at least this many iterations,
    /// but not less than `min_run_seconds`.
    pub fn iterations(&mut self, n: usize) -> &mut Self {
        self.iterations = n;
        self
    }

    /// Run measurement for at least `n` seconds (default: 3).
    /// The measurement will run for at least this many seconds, but not less than `iterations`.
    pub fn run_seconds(&mut self, n: usize) -> &mut Self {
        self.min_run_seconds = n;
        self
    }

    /// Warmup before measurement for min `n` seconds (default: 2).
    /// This is lower bound, and if iteration is still running after the warmup time,
    /// it will be allowed to finish before starting the timed iterations.
    pub fn warmup(&mut self, n: usize) -> &mut Self {
        self.warmup_seconds = n;
        self
    }

    /// Disable the automatic load / compare / print / save behaviour.
    ///
    /// With this set, `run()` only returns the [`BenchReport`] and performs no
    /// I/O or printing. Useful when calling bencher from tests or when you
    /// want to handle persistence yourself.
    pub fn no_auto_save(&mut self) -> &mut Self {
        self.auto_save = false;
        self
    }

    /// Set the benchmark name (default: fn name).
    pub fn name(&mut self, name: impl Into<String>) -> &mut Self {
        self.name = Some(name.into());
        self
    }

    /// Run the benchmark.
    ///
    /// Unless [`Bench::no_auto_save`] was called, this will:
    /// - Print a comparison if a previous baseline exists, otherwise print
    ///   the stats directly.
    /// - Save the result as the new baseline for the next run.
    pub fn run<F, R>(&mut self, f: F) -> BenchReport
    where
        F: FnMut() -> R,
    {
        let name = self.name.take().expect(
            "bencher: called run() twice on the same Bench, \
             call `name()` to create a new benchmark with the same params.",
        );
        let report = self.run_inner(&name, f);

        if self.auto_save {
            let path = self.baseline_path(&name);
            if let Some(parent) = path.parent() {
                std::fs::create_dir_all(parent).ok();
            }
            if let Ok(previous) = BenchReport::load(&path) {
                report.compare(&previous).print();
            } else {
                report.print();
            }
            report.save(&path).ok();
        }

        report
    }

    /// Path where the baseline JSON is stored/loaded.
    ///
    /// Resolves to `<CARGO_TARGET_DIR>/bencher/[<group>/]<name>/baseline.json`.
    /// Cargo sets `CARGO_TARGET_DIR` to the target directory when running benchmarks.
    pub fn baseline_path(&self, name: &str) -> PathBuf {
        let base = cargo_target_directory().unwrap();
        let mut path = base.join("bencher");
        if let Some(ref group) = self.group {
            path = path.join(group);
        }
        path.join(name).join("baseline.json")
    }

    // ── private ─────────────────────────────────────────────────────────────

    /// Warm up then collect timed iterations, returning the [`BenchReport`].
    pub(crate) fn run_inner<F, R>(&self, name: &str, mut f: F) -> BenchReport
    where
        F: FnMut() -> R,
    {
        assert!(self.iterations > 0, "bencher: iterations must be > 0");
        collector::init();

        // Warmup — discard spans
        let warmup_end =
            std::time::Instant::now() + std::time::Duration::from_secs(self.warmup_seconds as u64);
        while std::time::Instant::now() < warmup_end {
            run_iter(&mut f);
            collector::drain();
        }

        // Timed iterations: run until both the minimum iteration count and the
        // minimum run time are satisfied.
        let end_run =
            std::time::Instant::now() + std::time::Duration::from_secs(self.min_run_seconds as u64);
        let mut results = Vec::new();
        for iter in 0.. {
            run_iter(&mut f);
            results.push(collector::drain());

            if iter > self.iterations && std::time::Instant::now() >= end_run {
                break;
            }
        }

        BenchReport::from_iters(name.to_owned(), results)
    }
}

// ── BenchGroup ───────────────────────────────────────────────────────────────

/// Accumulates multiple benchmarks belonging to the same group and prints a
/// single combined table on drop.
///
/// Obtained via [`Bench::group`].
pub struct BenchGroup<'a> {
    bench: &'a mut Bench,
    group_name: String,
    entries: Vec<GroupEntry>,
}

struct GroupEntry {
    report: BenchReport,
    baseline: Option<BenchReport>,
    path: PathBuf,
}

impl BenchGroup<'_> {
    /// Set the name for the next benchmark in this group.
    pub fn name(&mut self, name: impl Into<String>) -> &mut Self {
        self.bench.name = Some(name.into());
        self
    }

    /// Number of timed iterations for subsequent benchmarks (default: 1000).
    pub fn iterations(&mut self, n: usize) -> &mut Self {
        self.bench.iterations = n;
        self
    }

    /// Run the next benchmark in this group.
    ///
    /// Results are collected and printed as a single combined table when the
    /// [`BenchGroup`] is dropped.
    pub fn run<F, R>(&mut self, f: F) -> BenchReport
    where
        F: FnMut() -> R,
    {
        let name = self
            .bench
            .name
            .take()
            .expect("bencher: call name() before each run() on a BenchGroup");
        let report = self.bench.run_inner(&name, f);
        let path = self.baseline_path(&name);
        let baseline = BenchReport::load(&path).ok();
        self.entries.push(GroupEntry {
            report: report.clone(),
            baseline,
            path,
        });
        report
    }

    fn baseline_path(&self, name: &str) -> PathBuf {
        let base = std::env::current_dir().unwrap_or_default();
        base.join("target")
            .join("bencher")
            .join(&self.group_name)
            .join(name)
            .join("baseline.json")
    }
}

impl Drop for BenchGroup<'_> {
    fn drop(&mut self) {
        if self.bench.auto_save {
            // Persist all reports first.
            for entry in &self.entries {
                if let Some(parent) = entry.path.parent() {
                    std::fs::create_dir_all(parent).ok();
                }
                entry.report.save(&entry.path).ok();
            }
            // Print combined group table.
            let pairs: Vec<(&BenchReport, Option<&BenchReport>)> = self
                .entries
                .iter()
                .map(|e| (&e.report, e.baseline.as_ref()))
                .collect();
            print_group(&pairs, &self.group_name);
        }
    }
}

// ── helpers ──────────────────────────────────────────────────────────────────
fn run_iter<F, R>(f: &mut F)
where
    F: FnMut() -> R,
{
    let root = Span::root(BENCH_ROOT, SpanContext::random());
    let _guard = root.set_local_parent();
    f();
    // _guard dropped first (restores thread-local), then root (completes trace)
    drop(_guard);
    drop(root);
}

fn cargo_target_directory() -> Option<PathBuf> {
    #[derive(serde::Deserialize)]
    struct Metadata {
        target_directory: PathBuf,
    }

    std::env::var_os("CARGO_TARGET_DIR")
        .map(PathBuf::from)
        .or_else(|| {
            let output =
                std::process::Command::new(std::env::var_os("CARGO").unwrap_or("cargo".into()))
                    .args(["metadata", "--format-version", "1"])
                    .output()
                    .ok()?;
            let metadata: Metadata = serde_json::from_slice(&output.stdout).ok()?;
            Some(metadata.target_directory)
        })
}