use core::sync::atomic::{AtomicU64, Ordering};

use super::page::page_size;

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct CompactionEstimate {
    pub candidate_heaps: u32,
    pub candidate_pages: u32,
    pub candidate_free_bytes: usize,
    pub best_case_meshes: u32,
    pub best_case_reclaimable_pages: u32,
    pub best_case_reclaimable_bytes: usize,
}

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct MeshStats {
    pub arena_size: usize,
    pub reserved_bytes: usize,
    pub reusable_span_count: u32,
    pub reusable_span_bytes: usize,
    pub live_miniheaps: u32,
    pub live_small_heaps: u32,
    pub partial_small_heaps: u32,
    pub full_small_heaps: u32,
    pub meshed_small_heaps: u32,
    pub reusable_small_heaps: u32,
    pub live_large_allocations: u32,
    pub live_small_bytes: usize,
    pub live_large_bytes: usize,
    pub retained_small_span_bytes: usize,
    pub retained_large_span_bytes: usize,
    pub virtual_small_span_bytes: usize,
    pub small_allocations: u64,
    pub small_deallocations: u64,
    pub large_allocations: u64,
    pub large_deallocations: u64,
    pub compact_calls: u64,
    pub meshes_performed: u64,
    pub meshed_pages: u64,
    pub meshed_bytes: u64,
    pub compaction: CompactionEstimate,
}

impl MeshStats {
    #[inline(always)]
    pub const fn live_bytes(&self) -> usize {
        self.live_small_bytes + self.live_large_bytes
    }

    #[inline(always)]
    pub const fn retained_bytes(&self) -> usize {
        self.retained_small_span_bytes + self.retained_large_span_bytes
    }

    #[inline(always)]
    pub const fn small_fragmentation_bytes(&self) -> usize {
        self.retained_small_span_bytes
            .saturating_sub(self.live_small_bytes)
    }

    #[inline(always)]
    pub const fn mesh_alias_bytes(&self) -> usize {
        self.virtual_small_span_bytes
            .saturating_sub(self.retained_small_span_bytes)
    }

    pub fn compaction_advice(&self) -> CompactionAdvice {
        let fragmented = self.small_fragmentation_bytes();
        let retained = self.retained_small_span_bytes;
        let fragmentation_percent = fragmented
            .saturating_mul(100)
            .checked_div(retained)
            .unwrap_or(0)
            .min(100) as u8;

        let reclaimable = self.compaction.best_case_reclaimable_bytes;
        let page = page_size();
        let recommendation = if self.compaction.best_case_meshes == 0 || reclaimable < page {
            CompactionRecommendation::Idle
        } else if reclaimable >= page * 4
            && (fragmentation_percent >= 20 || reclaimable.saturating_mul(4) >= retained.max(page))
        {
            CompactionRecommendation::Compact
        } else if reclaimable >= page
            && (fragmentation_percent >= 10 || self.compaction.candidate_heaps >= 2)
        {
            CompactionRecommendation::Consider
        } else {
            CompactionRecommendation::Idle
        };

        CompactionAdvice {
            recommendation,
            fragmentation_bytes: fragmented,
            fragmentation_percent,
            candidate_heaps: self.compaction.candidate_heaps,
            candidate_free_bytes: self.compaction.candidate_free_bytes,
            best_case_meshes: self.compaction.best_case_meshes,
            best_case_reclaimable_pages: self.compaction.best_case_reclaimable_pages,
            best_case_reclaimable_bytes: reclaimable,
        }
    }
}

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct CompactionAdvice {
    pub recommendation: CompactionRecommendation,
    pub fragmentation_bytes: usize,
    pub fragmentation_percent: u8,
    pub candidate_heaps: u32,
    pub candidate_free_bytes: usize,
    pub best_case_meshes: u32,
    pub best_case_reclaimable_pages: u32,
    pub best_case_reclaimable_bytes: usize,
}

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum CompactionRecommendation {
    #[default]
    Idle,
    Consider,
    Compact,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RuntimeCompactionPolicy {
    pub minimum_recommendation: CompactionRecommendation,
    pub min_fragmentation_bytes: usize,
    pub min_reclaimable_bytes: usize,
    pub min_candidate_heaps: u32,
}

impl RuntimeCompactionPolicy {
    pub fn should_compact(&self, advice: &CompactionAdvice) -> bool {
        recommendation_rank(advice.recommendation)
            >= recommendation_rank(self.minimum_recommendation)
            && advice.fragmentation_bytes >= self.min_fragmentation_bytes
            && advice.best_case_reclaimable_bytes >= self.min_reclaimable_bytes
            && advice.candidate_heaps >= self.min_candidate_heaps
    }
}

impl Default for RuntimeCompactionPolicy {
    fn default() -> Self {
        let page = page_size();
        Self {
            minimum_recommendation: CompactionRecommendation::Consider,
            min_fragmentation_bytes: page,
            min_reclaimable_bytes: page,
            min_candidate_heaps: 2,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RuntimeCompactionResult {
    Compacted {
        meshes: usize,
        advice: CompactionAdvice,
    },
    Skipped {
        reason: CompactionSkipReason,
        advice: Option<CompactionAdvice>,
    },
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CompactionSkipReason {
    AllocatorUnavailable,
    ThreadUnavailable,
    NotAtSafepoint,
    ThreadsActive,
    Policy,
}

#[inline(always)]
const fn recommendation_rank(recommendation: CompactionRecommendation) -> u8 {
    match recommendation {
        CompactionRecommendation::Idle => 0,
        CompactionRecommendation::Consider => 1,
        CompactionRecommendation::Compact => 2,
    }
}

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(crate) struct CounterSnapshot {
    pub small_allocations: u64,
    pub small_deallocations: u64,
    pub large_allocations: u64,
    pub large_deallocations: u64,
    pub compact_calls: u64,
    pub meshes_performed: u64,
    pub meshed_pages: u64,
    pub meshed_bytes: u64,
}

#[derive(Debug)]
pub(crate) struct StatsState {
    small_allocations: AtomicU64,
    small_deallocations: AtomicU64,
    large_allocations: AtomicU64,
    large_deallocations: AtomicU64,
    compact_calls: AtomicU64,
    meshes_performed: AtomicU64,
    meshed_pages: AtomicU64,
    meshed_bytes: AtomicU64,
}

impl StatsState {
    pub const fn new() -> Self {
        Self {
            small_allocations: AtomicU64::new(0),
            small_deallocations: AtomicU64::new(0),
            large_allocations: AtomicU64::new(0),
            large_deallocations: AtomicU64::new(0),
            compact_calls: AtomicU64::new(0),
            meshes_performed: AtomicU64::new(0),
            meshed_pages: AtomicU64::new(0),
            meshed_bytes: AtomicU64::new(0),
        }
    }

    #[inline(always)]
    pub fn record_small_allocation(&self) {
        self.small_allocations.fetch_add(1, Ordering::Relaxed);
    }

    #[inline(always)]
    pub fn record_small_deallocation(&self) {
        self.small_deallocations.fetch_add(1, Ordering::Relaxed);
    }

    #[inline(always)]
    pub fn record_large_allocation(&self) {
        self.large_allocations.fetch_add(1, Ordering::Relaxed);
    }

    #[inline(always)]
    pub fn record_large_deallocation(&self) {
        self.large_deallocations.fetch_add(1, Ordering::Relaxed);
    }

    #[inline(always)]
    pub fn record_compact_call(&self) {
        self.compact_calls.fetch_add(1, Ordering::Relaxed);
    }

    #[inline(always)]
    pub fn record_mesh(&self, pages: u32, bytes: usize) {
        self.meshes_performed.fetch_add(1, Ordering::Relaxed);
        self.meshed_pages.fetch_add(pages as u64, Ordering::Relaxed);
        self.meshed_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
    }

    pub fn snapshot(&self) -> CounterSnapshot {
        CounterSnapshot {
            small_allocations: self.small_allocations.load(Ordering::Relaxed),
            small_deallocations: self.small_deallocations.load(Ordering::Relaxed),
            large_allocations: self.large_allocations.load(Ordering::Relaxed),
            large_deallocations: self.large_deallocations.load(Ordering::Relaxed),
            compact_calls: self.compact_calls.load(Ordering::Relaxed),
            meshes_performed: self.meshes_performed.load(Ordering::Relaxed),
            meshed_pages: self.meshed_pages.load(Ordering::Relaxed),
            meshed_bytes: self.meshed_bytes.load(Ordering::Relaxed),
        }
    }
}
