//! deterministic heap allocation and memory footprint benchmarks for repo ingestion:
//! car parsing, commit proof parsing, mst node decoding, full mst walks, sparse scans,
//! and streaming filtered cars (in-order and reverse out-of-order).
//!
//! wraps the production allocator from `src/allocator.rs` in a counting `GlobalAlloc` so
//! every benchmark reports exact `allocs`, `reallocs`, `total_bytes` (churn), `peak_bytes`
//! (high-water mark above baseline), and `retained_bytes` (live bytes held in the result).

#[path = "../src/car.rs"]
#[allow(dead_code)]
mod car;
#[cfg(feature = "indexer")]
#[path = "../src/filter.rs"]
#[allow(dead_code, unused_imports)]
mod filter;
#[cfg(feature = "indexer")]
#[path = "../src/backfill/filtered_car.rs"]
#[allow(dead_code, unused_imports)]
mod filtered_car;
#[path = "../src/mst.rs"]
#[allow(dead_code)]
mod mst;
#[path = "../src/sparse_mst.rs"]
#[allow(dead_code, unused_imports)]
mod sparse_mst;

use core::alloc::{GlobalAlloc, Layout};
use std::collections::HashSet;
use std::hint::black_box;
use std::path::PathBuf;
#[cfg(feature = "indexer")]
use std::sync::Arc;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};

use bytes::Bytes;
use cid::Cid as IpldCid;
use smol_str::SmolStr;

use car::{CarBlock, CarBlocks};
use mst::{FlatEntry, MstKey};
use sparse_mst::{SparseScanner, sparse_ranges};

#[cfg(feature = "alloc-mimalloc")]
static INNER: mimalloc::MiMalloc = mimalloc::MiMalloc;

#[cfg(feature = "alloc-jemalloc")]
static INNER: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

#[cfg(all(feature = "alloc-jemalloc", target_os = "linux"))]
#[allow(non_upper_case_globals)]
#[unsafe(export_name = "_rjem_malloc_conf")]
pub static _rjem_malloc_conf: Option<&'static core::ffi::CStr> =
    Some(c"background_thread:true,percpu_arena:percpu,dirty_decay_ms:1000,muzzy_decay_ms:0");

#[cfg(feature = "alloc-snmalloc")]
static INNER: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;

#[cfg(not(any(
    feature = "alloc-mimalloc",
    feature = "alloc-jemalloc",
    feature = "alloc-snmalloc"
)))]
static INNER: std::alloc::System = std::alloc::System;

struct CountingAlloc;

#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;

static MEASURING: AtomicBool = AtomicBool::new(false);
static ALLOCS: AtomicU64 = AtomicU64::new(0);
static DEALLOCS: AtomicU64 = AtomicU64::new(0);
static REALLOCS: AtomicU64 = AtomicU64::new(0);
static TOTAL_BYTES: AtomicU64 = AtomicU64::new(0);
static LIVE_BYTES: AtomicI64 = AtomicI64::new(0);
static PEAK_BYTES: AtomicI64 = AtomicI64::new(0);

unsafe impl GlobalAlloc for CountingAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { INNER.alloc(layout) };
        if !ptr.is_null() && MEASURING.load(Ordering::Relaxed) {
            let size = layout.size() as u64;
            ALLOCS.fetch_add(1, Ordering::Relaxed);
            TOTAL_BYTES.fetch_add(size, Ordering::Relaxed);
            let live = LIVE_BYTES.fetch_add(size as i64, Ordering::Relaxed) + size as i64;
            PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
        }
        ptr
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { INNER.alloc_zeroed(layout) };
        if !ptr.is_null() && MEASURING.load(Ordering::Relaxed) {
            let size = layout.size() as u64;
            ALLOCS.fetch_add(1, Ordering::Relaxed);
            TOTAL_BYTES.fetch_add(size, Ordering::Relaxed);
            let live = LIVE_BYTES.fetch_add(size as i64, Ordering::Relaxed) + size as i64;
            PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
        }
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if MEASURING.load(Ordering::Relaxed) {
            DEALLOCS.fetch_add(1, Ordering::Relaxed);
            LIVE_BYTES.fetch_sub(layout.size() as i64, Ordering::Relaxed);
        }
        unsafe { INNER.dealloc(ptr, layout) };
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let new_ptr = unsafe { INNER.realloc(ptr, layout, new_size) };
        if !new_ptr.is_null() && MEASURING.load(Ordering::Relaxed) {
            REALLOCS.fetch_add(1, Ordering::Relaxed);
            let old_size = layout.size();
            if new_size > old_size {
                let grow = (new_size - old_size) as u64;
                TOTAL_BYTES.fetch_add(grow, Ordering::Relaxed);
                let live = LIVE_BYTES.fetch_add(grow as i64, Ordering::Relaxed) + grow as i64;
                PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
            } else {
                let shrink = (old_size - new_size) as i64;
                LIVE_BYTES.fetch_sub(shrink, Ordering::Relaxed);
            }
        }
        new_ptr
    }
}

#[derive(Clone, Copy)]
struct Stats {
    allocs: u64,
    deallocs: u64,
    reallocs: u64,
    total_bytes: u64,
    peak_bytes: u64,
    retained_bytes: u64,
}

fn measure<T>(f: impl FnOnce() -> T) -> Stats {
    ALLOCS.store(0, Ordering::SeqCst);
    DEALLOCS.store(0, Ordering::SeqCst);
    REALLOCS.store(0, Ordering::SeqCst);
    TOTAL_BYTES.store(0, Ordering::SeqCst);
    LIVE_BYTES.store(0, Ordering::SeqCst);
    PEAK_BYTES.store(0, Ordering::SeqCst);
    MEASURING.store(true, Ordering::SeqCst);

    let out = black_box(f());

    MEASURING.store(false, Ordering::SeqCst);
    let retained_bytes = LIVE_BYTES.load(Ordering::SeqCst).max(0) as u64;
    let peak_bytes = PEAK_BYTES.load(Ordering::SeqCst).max(0) as u64;
    let stats = Stats {
        allocs: ALLOCS.load(Ordering::SeqCst),
        deallocs: DEALLOCS.load(Ordering::SeqCst),
        reallocs: REALLOCS.load(Ordering::SeqCst),
        total_bytes: TOTAL_BYTES.load(Ordering::SeqCst),
        peak_bytes,
        retained_bytes,
    };
    drop(out);
    stats
}

const REPOS: [&str; 3] = ["atproto.com", "jay.bsky.team", "pfrazee.com"];
const PROOFS: usize = 2000;
const DENSE_COLLECTION: &str = "app.bsky.feed.post";
const SELECTIVE_COLLECTION: &str = "sh.tangled.*";
const CHECKS: [(&str, bool); 2] = [("verify", true), ("trust", false)];
const BODY_CHUNK: usize = 16 * 1024;

struct Repo {
    name: &'static str,
    car: Bytes,
    #[cfg(feature = "indexer")]
    reverse_car: OnceLock<Bytes>,
    mst_root: IpldCid,
    mst_nodes: Vec<(IpldCid, Bytes)>,
    leaves: Vec<(MstKey, IpldCid)>,
}

impl Repo {
    fn parse(&self, verify_cids: bool) -> car::ParsedCar<CarBlocks> {
        car::parse_car(self.car.clone(), verify_cids).expect("bench repo parses")
    }

    /// the same car with record blocks first, mst nodes in bottom-up post-order (children before
    /// parents), and the commit block last, exercising out-of-order streaming and node parking.
    #[cfg(feature = "indexer")]
    #[cfg(feature = "indexer")]
    fn reverse_car(&self) -> &Bytes {
        self.reverse_car.get_or_init(|| {
            let parsed = self.parse(false);
            let blocks = parsed.blocks.map();
            let mst_cids: HashSet<IpldCid> = self.mst_nodes.iter().map(|(cid, _)| *cid).collect();
            let mut records: Vec<(IpldCid, &Bytes)> = blocks
                .iter()
                .filter(|(cid, _)| **cid != parsed.root && !mst_cids.contains(cid))
                .map(|(cid, bytes)| (*cid, bytes))
                .collect();
            records.sort_unstable_by_key(|(cid, _)| *cid);
            let sections = records
                .into_iter()
                .chain(self.mst_nodes.iter().rev().map(|(cid, bytes)| (*cid, bytes)))
                .chain(std::iter::once((parsed.root, &blocks[&parsed.root])));
            write_car(parsed.root, sections)
        })
    }
}

fn data_dir() -> PathBuf {
    std::env::var_os("HYDRANT_BENCH_DATA")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("target/bench-data"))
}

fn repos() -> &'static [Repo] {
    static LOADED: OnceLock<Vec<Repo>> = OnceLock::new();
    LOADED.get_or_init(|| REPOS.into_iter().map(load).collect())
}

fn load(name: &'static str) -> Repo {
    let path = data_dir().join(format!("{name}.car"));
    let car = std::fs::read(&path).unwrap_or_else(|e| {
        panic!(
            "{}: {e}; fetch the bench repos with `nu tests/bench.nu fetch`",
            path.display()
        )
    });
    let car = Bytes::from(car);
    let parsed = car::parse_car(car.clone(), true).expect("bench repo is a valid car");
    let mst_root = commit_data(&parsed);
    let mst_nodes = collect_mst_nodes(parsed.blocks.map(), mst_root);
    let leaves = mst::leaves(parsed.blocks.map(), mst_root).expect("bench repo mst walks");
    Repo {
        name,
        car,
        #[cfg(feature = "indexer")]
        reverse_car: OnceLock::new(),
        mst_root,
        mst_nodes,
        leaves,
    }
}

fn collect_mst_nodes(blocks: &CarBlocks, root: IpldCid) -> Vec<(IpldCid, Bytes)> {
    let mut out = Vec::new();
    let mut visited = HashSet::new();
    let mut stack = vec![root];
    while let Some(cid) = stack.pop() {
        if !visited.insert(cid) {
            continue;
        }
        let bytes = blocks.get(&cid).expect("mst node in car").clone();
        let entries = mst::decode_node(&bytes).expect("bench repo mst node decodes");
        let start = stack.len();
        for entry in entries {
            if let FlatEntry::Tree { cid } = entry {
                stack.push(cid);
            }
        }
        stack[start..].reverse();
        out.push((cid, bytes));
    }
    out
}

fn commit_data(parsed: &car::ParsedCar<CarBlocks>) -> IpldCid {
    let commit = parsed.blocks.get(&parsed.root).expect("car has its commit");
    mst::AtpCommit::from_cbor(commit)
        .expect("commit decodes")
        .data
}

fn proofs() -> &'static [Bytes] {
    static CUT: OnceLock<Vec<Bytes>> = OnceLock::new();
    CUT.get_or_init(|| {
        let repo = repos().last().expect("bench repos are listed");
        let parsed = repo.parse(false);
        let blocks = parsed.blocks.map();
        let step = (repo.leaves.len() / PROOFS).max(1);
        repo.leaves
            .iter()
            .step_by(step)
            .take(PROOFS)
            .map(|(key, record)| {
                let cids = std::iter::once(parsed.root)
                    .chain(mst_path(blocks, repo.mst_root, key))
                    .chain([*record]);
                let proof = write_car(parsed.root, cids.map(|cid| (cid, &blocks[&cid])));
                car::parse_commit_car(proof.clone(), true).expect("proof is a valid car");
                proof
            })
            .collect()
    })
}

fn mst_path(blocks: &CarBlocks, root: IpldCid, key: &str) -> Vec<IpldCid> {
    let mut path = vec![root];
    loop {
        let node = path.last().expect("path starts at the root");
        let entries = mst::decode_node(&blocks[node]).expect("bench repo nodes decode");
        let mut subtree = None;
        for entry in &entries {
            match entry {
                FlatEntry::Leaf { key: leaf, .. } if leaf == key => return path,
                FlatEntry::Leaf { key: leaf, .. } if leaf.as_str() > key => break,
                FlatEntry::Leaf { .. } => subtree = None,
                FlatEntry::Tree { cid } => subtree = Some(*cid),
            }
        }
        path.push(subtree.expect("key is in the tree"));
    }
}

fn write_car<'a>(root: IpldCid, sections: impl Iterator<Item = (IpldCid, &'a Bytes)>) -> Bytes {
    let header = car::CarHeader::new_v1(vec![root])
        .encode()
        .expect("header encodes");
    let mut out = Vec::new();
    put_uvarint(&mut out, header.len());
    out.extend_from_slice(&header);
    for (cid, bytes) in sections {
        let cid = cid.to_bytes();
        put_uvarint(&mut out, cid.len() + bytes.len());
        out.extend_from_slice(&cid);
        out.extend_from_slice(bytes);
    }
    out.into()
}

fn put_uvarint(out: &mut Vec<u8>, mut value: usize) {
    while value >= 0x80 {
        out.push(value as u8 | 0x80);
        value >>= 7;
    }
    out.push(value as u8);
}

#[cfg(feature = "indexer")]
fn filter_for(pattern: &str) -> Arc<filter::FilterConfig> {
    let mut cfg = filter::FilterConfig::new(filter::FilterMode::Filter);
    cfg.collections.push(SmolStr::new(pattern));
    Arc::new(cfg)
}

#[cfg(feature = "indexer")]
fn run_filtered_car(
    car_bytes: &Bytes,
    mst_root: IpldCid,
    filter: &Arc<filter::FilterConfig>,
) -> (
    car::Blocks<car::CommitBlocks>,
    sparse_mst::SparseScanOutput,
) {
    let mut filtered = filtered_car::FilteredCar::new(filter.clone(), false).unwrap();
    for chunk in car_bytes.chunks(BODY_CHUNK) {
        filtered.push(chunk).unwrap();
    }
    let mut repo = filtered.finish().unwrap();
    let scan = repo.scanner.scan(mst_root).unwrap().unwrap();
    (repo.blocks, scan)
}

fn record(
    rows: &mut Vec<(String, Stats)>,
    filter: Option<&str>,
    name: impl Into<String>,
    run: impl FnOnce() -> Stats,
) {
    let name = name.into();
    if filter.is_some_and(|pat| !name.contains(pat)) {
        return;
    }
    rows.push((name, run()));
}

fn main() {
    let mut json = false;
    let mut filter: Option<String> = None;
    for arg in std::env::args().skip(1) {
        match arg.as_str() {
            "--bench" => {}
            "--json" => json = true,
            other if !other.starts_with('-') => filter = Some(other.to_string()),
            _ => {}
        }
    }
    let filter = filter.as_deref();
    let mut rows = Vec::new();

    for (check, verify_cids) in CHECKS {
        for repo in repos() {
            let car = repo.car.clone();
            record(
                &mut rows,
                filter,
                format!("parse_car/{check}/{}", repo.name),
                || measure(|| car::parse_car(car, verify_cids).unwrap()),
            );
        }
    }

    let proofs = proofs();
    for (check, verify_cids) in CHECKS {
        record(&mut rows, filter, format!("parse_commit_car/{check}"), || {
            measure(|| {
                for proof in proofs {
                    black_box(car::parse_commit_car(proof.clone(), verify_cids).unwrap());
                }
            })
        });
    }

    for repo in repos() {
        record(&mut rows, filter, format!("mst_decode/{}", repo.name), || {
            measure(|| {
                for (_, bytes) in &repo.mst_nodes {
                    black_box(mst::decode_node(bytes).unwrap());
                }
            })
        });
    }

    for repo in repos() {
        let parsed = repo.parse(false);
        record(&mut rows, filter, format!("mst_walk/{}", repo.name), || {
            measure(|| {
                let mut count = 0usize;
                for leaf in mst::walk(parsed.blocks.map(), repo.mst_root) {
                    black_box(leaf.unwrap());
                    count += 1;
                }
                count
            })
        });
    }

    for repo in repos() {
        let parsed = repo.parse(false);
        record(&mut rows, filter, format!("mst_leaves/{}", repo.name), || {
            measure(|| mst::leaves(parsed.blocks.map(), repo.mst_root).unwrap())
        });
    }

    for (label, pattern) in [
        ("dense", DENSE_COLLECTION),
        ("selective", SELECTIVE_COLLECTION),
    ] {
        let ranges = sparse_ranges(&[SmolStr::new(pattern)]);
        for repo in repos() {
            let blocks = car::parse_car_blocks(repo.car.clone(), false).unwrap();
            let ranges = ranges.clone();
            record(
                &mut rows,
                filter,
                format!("sparse_scan/{label}/{}", repo.name),
                move || {
                    measure(move || {
                        let mut scanner = SparseScanner::new(ranges, blocks);
                        let scan = scanner.scan(repo.mst_root).unwrap().unwrap();
                        (scanner, scan)
                    })
                },
            );
        }
    }

    #[cfg(feature = "indexer")]
    {
        let selective = filter_for(SELECTIVE_COLLECTION);
        let dense = filter_for(DENSE_COLLECTION);
        for (label, cfg) in [("selective", &selective), ("dense", &dense)] {
            for repo in repos() {
                record(
                    &mut rows,
                    filter,
                    format!("filtered_car/{label}/{}", repo.name),
                    || measure(|| run_filtered_car(&repo.car, repo.mst_root, cfg)),
                );
            }
        }
        for repo in repos() {
            let reversed = repo.reverse_car();
            record(
                &mut rows,
                filter,
                format!("filtered_car/reverse/{}", repo.name),
                || measure(|| run_filtered_car(reversed, repo.mst_root, &selective)),
            );
        }
    }

    for repo in repos() {
        for (check, verify_cids) in CHECKS {
            let parsed = repo.parse(verify_cids);
            let records: Vec<CarBlock> = repo
                .leaves
                .iter()
                .map(|(_, cid)| parsed.blocks.block(cid).expect("record is in the car"))
                .collect();
            record(
                &mut rows,
                filter,
                format!("record_check/{check}/{}", repo.name),
                || {
                    measure(|| {
                        for block in &records {
                            black_box(block.verify().unwrap());
                        }
                    })
                },
            );
        }
    }

    for (check, verify_cids) in CHECKS {
        for repo in repos() {
            record(
                &mut rows,
                filter,
                format!("backfill_cpu/{check}/{}", repo.name),
                || {
                    measure(|| {
                        let mut whole = car::WholeCar::new(verify_cids);
                        whole.reserve(repo.car.len());
                        for chunk in repo.car.chunks(BODY_CHUNK) {
                            whole.push(chunk).unwrap();
                        }
                        let parsed = whole.finish().unwrap();
                        let leaves =
                            mst::leaves(parsed.blocks.map(), commit_data(&parsed)).unwrap();
                        for (_, cid) in &leaves {
                            black_box(parsed.blocks.block(cid).unwrap().verify().unwrap());
                        }
                        (parsed, leaves)
                    })
                },
            );
        }
    }

    if json {
        let items: Vec<serde_json::Value> = rows
            .into_iter()
            .map(|(bench, s)| {
                serde_json::json!({
                    "bench": bench,
                    "allocs": s.allocs,
                    "deallocs": s.deallocs,
                    "reallocs": s.reallocs,
                    "total_bytes": s.total_bytes,
                    "peak_bytes": s.peak_bytes,
                    "retained_bytes": s.retained_bytes,
                })
            })
            .collect();
        println!("{}", serde_json::to_string_pretty(&items).unwrap());
    } else {
        println!(
            "{:<38} {:>10} {:>10} {:>14} {:>14} {:>14}",
            "bench", "allocs", "reallocs", "total_bytes", "peak_bytes", "retained_bytes"
        );
        println!("{}", "-".repeat(104));
        for (bench, s) in rows {
            println!(
                "{:<38} {:>10} {:>10} {:>14} {:>14} {:>14}",
                bench, s.allocs, s.reallocs, s.total_bytes, s.peak_bytes, s.retained_bytes
            );
        }
    }
}