use core::alloc::{GlobalAlloc, Layout};
use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::ptr::{addr_of_mut, drop_in_place, null, null_mut};
use core::sync::atomic::{AtomicU32, Ordering};

use super::allocator::MeshAllocator;
use super::constants::DEFAULT_ARENA_SIZE;
use super::stats::{
    CompactionAdvice, CompactionSkipReason, MeshStats, RuntimeCompactionPolicy,
    RuntimeCompactionResult,
};
use super::sync::{FutexMutex, futex_wait_for_value, futex_wake_all};
use super::thread_local_heap::ThreadLocalHeap;

const INIT_UNINITIALIZED: u32 = 0;
const INIT_IN_PROGRESS: u32 = 1;
const INIT_READY: u32 = 2;
const INIT_FAILED: u32 = 3;

pub const DEFAULT_GLOBAL_MINIHEAP_CAPACITY: u32 = 4096;

const TLS_UNINITIALIZED: u32 = 0;
const TLS_READY: u32 = 1;
const TLS_FAILED: u32 = 2;
const SAFEPOINT_INACTIVE: u32 = 0;
const SAFEPOINT_ACTIVE: u32 = 1;

#[thread_local]
static mut THREAD_HEAP_STATE: u32 = TLS_UNINITIALIZED;
#[thread_local]
static mut THREAD_HEAP: MaybeUninit<ThreadLocalHeap> = MaybeUninit::uninit();
#[thread_local]
static mut THREAD_SAFEPOINT_STATE: u32 = SAFEPOINT_INACTIVE;
#[thread_local]
static mut THREAD_HEAP_OWNER: *const GlobalMeshAllocator = null();

#[derive(Debug)]
pub struct GlobalMeshAllocator {
    arena_size: usize,
    miniheap_capacity: u32,
    init_state: AtomicU32,
    registered_threads: AtomicU32,
    quiescent_threads: AtomicU32,
    lock: FutexMutex,
    allocator: UnsafeCell<MaybeUninit<MeshAllocator>>,
}

impl GlobalMeshAllocator {
    pub const fn new(arena_size: usize, miniheap_capacity: u32) -> Self {
        Self {
            arena_size,
            miniheap_capacity,
            init_state: AtomicU32::new(INIT_UNINITIALIZED),
            registered_threads: AtomicU32::new(0),
            quiescent_threads: AtomicU32::new(0),
            lock: FutexMutex::new(),
            allocator: UnsafeCell::new(MaybeUninit::uninit()),
        }
    }

    pub const fn with_default_config() -> Self {
        Self::new(DEFAULT_ARENA_SIZE, DEFAULT_GLOBAL_MINIHEAP_CAPACITY)
    }

    pub fn init_thread(&self) -> bool {
        self.thread_heap().is_some()
    }

    pub fn shutdown_thread(&self) {
        unsafe {
            let state_ptr = addr_of_mut!(THREAD_HEAP_STATE);
            let heap_ptr = addr_of_mut!(THREAD_HEAP);
            let safepoint_ptr = addr_of_mut!(THREAD_SAFEPOINT_STATE);
            let owner_ptr = addr_of_mut!(THREAD_HEAP_OWNER);
            if *safepoint_ptr == SAFEPOINT_ACTIVE {
                self.quiescent_threads.fetch_sub(1, Ordering::AcqRel);
                *safepoint_ptr = SAFEPOINT_INACTIVE;
            }
            if *state_ptr == TLS_READY {
                let _ = self.with_existing_allocator_mut(|allocator| {
                    allocator.shutdown_thread((*heap_ptr).assume_init_mut());
                });
                drop_in_place((*heap_ptr).as_mut_ptr());
                self.registered_threads.fetch_sub(1, Ordering::AcqRel);
            }
            *state_ptr = TLS_UNINITIALIZED;
            *owner_ptr = null();
        }
    }

    /// Marks the current thread quiescent for cooperative compaction.
    ///
    /// This is **not** true concurrent compaction: arbitrary loads/stores through existing raw
    /// pointers remain outside allocator control, so compaction is only safe once all registered
    /// allocator threads have voluntarily entered this state.
    pub fn enter_quiescent_compaction_state(&self) -> bool {
        if self.thread_heap().is_none() {
            return false;
        }

        unsafe {
            let safepoint_ptr = addr_of_mut!(THREAD_SAFEPOINT_STATE);
            if *safepoint_ptr == SAFEPOINT_ACTIVE {
                return true;
            }
            *safepoint_ptr = SAFEPOINT_ACTIVE;
        }
        self.quiescent_threads.fetch_add(1, Ordering::AcqRel);
        true
    }

    /// Leaves the cooperative quiescent compaction state for the current thread.
    pub fn leave_quiescent_compaction_state(&self) {
        unsafe {
            let safepoint_ptr = addr_of_mut!(THREAD_SAFEPOINT_STATE);
            if *safepoint_ptr == SAFEPOINT_ACTIVE {
                *safepoint_ptr = SAFEPOINT_INACTIVE;
                self.quiescent_threads.fetch_sub(1, Ordering::AcqRel);
            }
        }
    }

    /// Returns true only when every registered allocator thread is quiescent.
    ///
    /// This is a cooperative global-quiescence check, not a proof that active mutators can safely
    /// race with remap/migration.
    pub fn quiescent_compaction_ready(&self) -> bool {
        let registered = self.registered_threads.load(Ordering::Acquire);
        registered != 0 && registered == self.quiescent_threads.load(Ordering::Acquire)
    }

    /// Runs cooperative quiescent compaction when the current thread and every other registered
    /// allocator thread have voluntarily stopped allocator-visible activity.
    ///
    /// This API intentionally does **not** claim to provide concurrent compaction with active
    /// mutators. Achieving that would require heavier machinery such as page-fault mediation,
    /// signal handling, syscall retry/interposition, or equivalent runtime coordination for raw
    /// pointer accesses and kernel I/O into moving pages.
    pub fn compact_when_quiescent(
        &self,
        policy: RuntimeCompactionPolicy,
    ) -> RuntimeCompactionResult {
        if !self.current_thread_in_safepoint() {
            return RuntimeCompactionResult::Skipped {
                reason: CompactionSkipReason::NotAtSafepoint,
                advice: self.compaction_advice(),
            };
        }
        if !self.quiescent_compaction_ready() {
            return RuntimeCompactionResult::Skipped {
                reason: CompactionSkipReason::ThreadsActive,
                advice: self.compaction_advice(),
            };
        }

        let Some(thread_heap) = self.thread_heap() else {
            return RuntimeCompactionResult::Skipped {
                reason: CompactionSkipReason::ThreadUnavailable,
                advice: self.compaction_advice(),
            };
        };
        let Some(result) = self.with_existing_allocator_mut(|allocator| {
            allocator.shutdown_thread(thread_heap);
            let advice = allocator.stats().compaction_advice();
            if !policy.should_compact(&advice) {
                return RuntimeCompactionResult::Skipped {
                    reason: CompactionSkipReason::Policy,
                    advice: Some(advice),
                };
            }

            let meshes = allocator.compact_with_thread(thread_heap);
            RuntimeCompactionResult::Compacted { meshes, advice }
        }) else {
            return RuntimeCompactionResult::Skipped {
                reason: CompactionSkipReason::AllocatorUnavailable,
                advice: None,
            };
        };

        result
    }

    /// Compatibility alias for `enter_quiescent_compaction_state`.
    pub fn enter_compaction_safepoint(&self) -> bool {
        self.enter_quiescent_compaction_state()
    }

    /// Compatibility alias for `leave_quiescent_compaction_state`.
    pub fn leave_compaction_safepoint(&self) {
        self.leave_quiescent_compaction_state();
    }

    /// Compatibility alias for `quiescent_compaction_ready`.
    pub fn compaction_safepoint_ready(&self) -> bool {
        self.quiescent_compaction_ready()
    }

    /// Compatibility alias for `compact_when_quiescent`.
    pub fn compact_at_safepoint(&self, policy: RuntimeCompactionPolicy) -> RuntimeCompactionResult {
        self.compact_when_quiescent(policy)
    }

    pub fn compact(&self) -> usize {
        let Some(thread_heap) = self.thread_heap() else {
            return 0;
        };
        self.with_allocator(|allocator| allocator.compact_with_thread(thread_heap))
            .unwrap_or_default()
    }

    pub fn stats(&self) -> Option<MeshStats> {
        self.with_existing_allocator(|allocator| allocator.stats())
    }

    pub fn compaction_advice(&self) -> Option<CompactionAdvice> {
        self.stats().map(|stats| stats.compaction_advice())
    }

    fn with_allocator<R>(&self, f: impl FnOnce(&mut MeshAllocator) -> R) -> Option<R> {
        if !self.ensure_initialized() {
            return None;
        }

        let _guard = self.lock.lock();
        let allocator = unsafe { (&mut *self.allocator.get()).assume_init_mut() };
        Some(f(allocator))
    }

    fn with_existing_allocator<R>(&self, f: impl FnOnce(&MeshAllocator) -> R) -> Option<R> {
        if self.init_state.load(Ordering::Acquire) != INIT_READY {
            return None;
        }

        let _guard = self.lock.lock();
        if self.init_state.load(Ordering::Acquire) != INIT_READY {
            return None;
        }
        let allocator = unsafe { (&*self.allocator.get()).assume_init_ref() };
        Some(f(allocator))
    }

    fn with_existing_allocator_mut<R>(&self, f: impl FnOnce(&mut MeshAllocator) -> R) -> Option<R> {
        if self.init_state.load(Ordering::Acquire) != INIT_READY {
            return None;
        }

        let _guard = self.lock.lock();
        if self.init_state.load(Ordering::Acquire) != INIT_READY {
            return None;
        }
        let allocator = unsafe { (&mut *self.allocator.get()).assume_init_mut() };
        Some(f(allocator))
    }

    fn ensure_initialized(&self) -> bool {
        loop {
            match self.init_state.load(Ordering::Acquire) {
                INIT_READY => return true,
                INIT_FAILED => return false,
                INIT_UNINITIALIZED => {
                    if self
                        .init_state
                        .compare_exchange(
                            INIT_UNINITIALIZED,
                            INIT_IN_PROGRESS,
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        let result = MeshAllocator::new(self.arena_size, self.miniheap_capacity);
                        match result {
                            Ok(allocator) => unsafe {
                                (*self.allocator.get()).write(allocator);
                                self.init_state.store(INIT_READY, Ordering::Release);
                                futex_wake_all(&self.init_state);
                                return true;
                            },
                            Err(_) => {
                                self.init_state.store(INIT_FAILED, Ordering::Release);
                                futex_wake_all(&self.init_state);
                                return false;
                            }
                        }
                    }
                }
                INIT_IN_PROGRESS => futex_wait_for_value(&self.init_state, INIT_IN_PROGRESS),
                _ => return false,
            }
        }
    }

    fn thread_heap(&self) -> Option<&'static mut ThreadLocalHeap> {
        unsafe {
            let owner_ptr = addr_of_mut!(THREAD_HEAP_OWNER);
            let heap_ptr = addr_of_mut!(THREAD_HEAP);
            match THREAD_HEAP_STATE {
                TLS_READY if core::ptr::eq(*owner_ptr, self) => Some((*heap_ptr).assume_init_mut()),
                TLS_READY => {
                    self.reset_foreign_thread_heap();
                    self.thread_heap()
                }
                TLS_FAILED => None,
                TLS_UNINITIALIZED => match ThreadLocalHeap::new() {
                    Ok(heap) => {
                        (*heap_ptr).write(heap);
                        *addr_of_mut!(THREAD_SAFEPOINT_STATE) = SAFEPOINT_INACTIVE;
                        *owner_ptr = self as *const _;
                        self.registered_threads.fetch_add(1, Ordering::AcqRel);
                        THREAD_HEAP_STATE = TLS_READY;
                        Some((*heap_ptr).assume_init_mut())
                    }
                    Err(_) => {
                        THREAD_HEAP_STATE = TLS_FAILED;
                        None
                    }
                },
                _ => None,
            }
        }
    }

    fn current_thread_in_safepoint(&self) -> bool {
        let _ = self;
        unsafe {
            core::ptr::eq(THREAD_HEAP_OWNER, self) && THREAD_SAFEPOINT_STATE == SAFEPOINT_ACTIVE
        }
    }

    fn reset_foreign_thread_heap(&self) {
        let _ = self;
        unsafe {
            if THREAD_HEAP_STATE == TLS_READY {
                drop_in_place(addr_of_mut!(THREAD_HEAP).cast::<ThreadLocalHeap>());
            }
            THREAD_HEAP_STATE = TLS_UNINITIALIZED;
            THREAD_SAFEPOINT_STATE = SAFEPOINT_INACTIVE;
            THREAD_HEAP_OWNER = null();
        }
    }
}

impl Drop for GlobalMeshAllocator {
    fn drop(&mut self) {
        unsafe {
            if core::ptr::eq(THREAD_HEAP_OWNER, self) {
                self.shutdown_thread();
            }
        }

        if self.init_state.load(Ordering::Acquire) == INIT_READY {
            let _guard = self.lock.lock();
            if self.init_state.load(Ordering::Acquire) == INIT_READY {
                unsafe {
                    drop_in_place((&mut *self.allocator.get()).as_mut_ptr());
                }
                self.init_state.store(INIT_FAILED, Ordering::Release);
            }
        }
    }
}

unsafe impl Sync for GlobalMeshAllocator {}
unsafe impl Send for GlobalMeshAllocator {}

unsafe impl GlobalAlloc for GlobalMeshAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if !self.ensure_initialized() {
            return null_mut();
        }
        let Some(thread_heap) = self.thread_heap() else {
            return null_mut();
        };

        let allocator_ref = unsafe { (&*self.allocator.get()).assume_init_ref() };
        if let Some(class) = size_class_for_layout(layout)
            && let Some(ptr) = allocator_ref.try_allocate_small_local(thread_heap, class)
        {
            return ptr;
        }

        self.with_allocator(|allocator| {
            allocator
                .allocate_layout_with_thread(thread_heap, layout)
                .unwrap_or(null_mut())
        })
        .unwrap_or(null_mut())
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        if !self.ensure_initialized() {
            return;
        }
        let Some(thread_heap) = self.thread_heap() else {
            return;
        };

        let allocator_ref = unsafe { (&*self.allocator.get()).assume_init_ref() };
        if allocator_ref.try_deallocate_local(ptr, thread_heap) {
            return;
        }

        let _ = self.with_allocator(|allocator| allocator.deallocate_with_thread(ptr, thread_heap));
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ptr = unsafe { self.alloc(layout) };
        if !ptr.is_null() {
            unsafe {
                ptr.write_bytes(0, layout.size());
            }
        }
        ptr
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if !self.ensure_initialized() {
            return null_mut();
        }
        let Some(thread_heap) = self.thread_heap() else {
            return null_mut();
        };

        self.with_allocator(|allocator| unsafe {
            allocator
                .reallocate_with_thread(ptr, layout, new_size, thread_heap)
                .unwrap_or(null_mut())
        })
        .unwrap_or(null_mut())
    }
}

fn size_class_for_layout(layout: Layout) -> Option<u8> {
    if layout.align() > super::page::page_size() {
        return None;
    }
    let aligned_size = layout
        .size()
        .max(1)
        .checked_add(layout.align() - 1)
        .map(|value| value & !(layout.align() - 1))?;
    if aligned_size > super::constants::MAX_SMALL_ALLOCATION {
        return None;
    }
    let class = super::size_map::size_class_for(aligned_size)?;
    if super::size_map::byte_size_for_class(class).is_multiple_of(layout.align()) {
        Some(class)
    } else {
        None
    }
}
