use std::{
    alloc::{GlobalAlloc, Layout, System},
    hint::black_box,
    sync::atomic::{AtomicUsize, Ordering},
    time::{Duration, Instant},
};

use crank_metrics::{CacheOutcome, record_cache_outcome};
use metrics_util::debugging::DebuggingRecorder;

const SAMPLES: usize = 40;
const WORK_UNITS: u64 = 500_000;
const MAX_PROFILE_ALLOCATION_BYTES: usize = 64 * 1024;

struct TrackingAllocator;

static CURRENT_HEAP_BYTES: AtomicUsize = AtomicUsize::new(0);
static PEAK_HEAP_BYTES: AtomicUsize = AtomicUsize::new(0);
static TOTAL_ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);

#[global_allocator]
static ALLOCATOR: TrackingAllocator = TrackingAllocator;

unsafe impl GlobalAlloc for TrackingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let pointer = unsafe { System.alloc(layout) };
        if !pointer.is_null() {
            record_allocation(layout.size());
        }
        pointer
    }

    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
        unsafe { System.dealloc(pointer, layout) };
        CURRENT_HEAP_BYTES.fetch_sub(layout.size(), Ordering::Relaxed);
    }

    unsafe fn realloc(&self, pointer: *mut u8, old: Layout, new_size: usize) -> *mut u8 {
        let resized = unsafe { System.realloc(pointer, old, new_size) };
        if !resized.is_null() {
            match new_size.cmp(&old.size()) {
                std::cmp::Ordering::Greater => record_allocation(new_size - old.size()),
                std::cmp::Ordering::Less => {
                    CURRENT_HEAP_BYTES.fetch_sub(old.size() - new_size, Ordering::Relaxed);
                }
                std::cmp::Ordering::Equal => {}
            }
        }
        resized
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let pointer = unsafe { System.alloc_zeroed(layout) };
        if !pointer.is_null() {
            record_allocation(layout.size());
        }
        pointer
    }
}

fn record_allocation(bytes: usize) {
    TOTAL_ALLOCATED_BYTES.fetch_add(bytes, Ordering::Relaxed);
    let current = CURRENT_HEAP_BYTES.fetch_add(bytes, Ordering::Relaxed) + bytes;
    PEAK_HEAP_BYTES.fetch_max(current, Ordering::Relaxed);
}

#[derive(Clone, Copy)]
struct ProfileSample {
    elapsed: Duration,
    allocated_bytes: usize,
    peak_heap_delta: usize,
}

#[test]
fn reproducible_foundation_profile_stays_inside_latency_and_cpu_proxy_budgets() {
    let recorder = DebuggingRecorder::new();
    let snapshotter = recorder.snapshotter();
    let mut baseline = Vec::with_capacity(SAMPLES);
    let mut enabled = Vec::with_capacity(SAMPLES);

    metrics::with_local_recorder(&recorder, || {
        record_cache_outcome(CacheOutcome::Hit);
        for index in 0..SAMPLES {
            if index % 2 == 0 {
                baseline.push(sample(false));
                enabled.push(sample(true));
            } else {
                enabled.push(sample(true));
                baseline.push(sample(false));
            }
        }
    });

    baseline.sort_unstable_by_key(|sample| sample.elapsed);
    enabled.sort_unstable_by_key(|sample| sample.elapsed);
    let baseline_p95 = baseline[SAMPLES * 95 / 100].elapsed;
    let enabled_p95 = enabled[SAMPLES * 95 / 100].elapsed;
    let baseline_cpu: Duration = baseline.iter().map(|sample| sample.elapsed).sum();
    let enabled_cpu: Duration = enabled.iter().map(|sample| sample.elapsed).sum();
    let baseline_allocated: usize = baseline.iter().map(|sample| sample.allocated_bytes).sum();
    let enabled_allocated: usize = enabled.iter().map(|sample| sample.allocated_bytes).sum();
    let baseline_peak = baseline
        .iter()
        .map(|sample| sample.peak_heap_delta)
        .max()
        .unwrap_or_default();
    let enabled_peak = enabled
        .iter()
        .map(|sample| sample.peak_heap_delta)
        .max()
        .unwrap_or_default();
    let series = snapshotter.snapshot().into_vec().len();

    println!(
        "metrics_profile_v1 samples={SAMPLES} work_units={WORK_UNITS} baseline_p95_ns={} enabled_p95_ns={} baseline_cpu_ns={} enabled_cpu_ns={} baseline_allocated_bytes={baseline_allocated} enabled_allocated_bytes={enabled_allocated} baseline_peak_heap_delta={baseline_peak} enabled_peak_heap_delta={enabled_peak} logical_series={series}",
        baseline_p95.as_nanos(),
        enabled_p95.as_nanos(),
        baseline_cpu.as_nanos(),
        enabled_cpu.as_nanos()
    );
    assert!(enabled_p95.as_nanos() * 100 <= baseline_p95.as_nanos() * 105);
    assert!(enabled_cpu.as_nanos() * 100 <= baseline_cpu.as_nanos() * 110);
    assert!(enabled_allocated <= baseline_allocated + MAX_PROFILE_ALLOCATION_BYTES);
    assert!(enabled_peak <= baseline_peak + MAX_PROFILE_ALLOCATION_BYTES);
    assert_eq!(series, 1);
}

fn sample(with_metrics: bool) -> ProfileSample {
    let heap_before = CURRENT_HEAP_BYTES.load(Ordering::Relaxed);
    let allocated_before = TOTAL_ALLOCATED_BYTES.load(Ordering::Relaxed);
    PEAK_HEAP_BYTES.store(heap_before, Ordering::Relaxed);
    let started = Instant::now();
    let mut value = 0x9e37_79b9_u64;
    for index in 0..WORK_UNITS {
        value = value.rotate_left(7) ^ index.wrapping_mul(0x100_0000_01b3);
    }
    black_box(value);
    if with_metrics {
        record_cache_outcome(CacheOutcome::Hit);
    }
    let elapsed = started.elapsed();
    let allocated_bytes = TOTAL_ALLOCATED_BYTES
        .load(Ordering::Relaxed)
        .saturating_sub(allocated_before);
    let peak_heap_delta = PEAK_HEAP_BYTES
        .load(Ordering::Relaxed)
        .saturating_sub(heap_before);
    ProfileSample {
        elapsed,
        allocated_bytes,
        peak_heap_delta,
    }
}