use core::cell::UnsafeCell;
use core::mem::size_of;
use core::ptr::null_mut;
use core::sync::atomic::{AtomicU32, Ordering};

use super::constants::DEFAULT_ARENA_SIZE;
use super::miniheap::MiniHeapId;
use super::page::{PageConfig, page_count, round_up_to_page};
use super::platform;
use super::raw_sys;
use super::span::Span;
use super::sync::FutexMutex;

const MAX_FREE_SPANS: usize = 4096;

#[derive(Debug)]
pub struct Arena {
    config: PageConfig,
    arena_size: usize,
    page_count: u32,
    fd: i32,
    base: *mut u8,
    owners: *mut AtomicU32,
    next_page: AtomicU32,
    free_spans: *mut Span,
    free_span_count: UnsafeCell<u32>,
    free_span_lock: FutexMutex,
}

impl Arena {
    #[inline]
    pub fn new() -> raw_sys::Result<Self> {
        Self::with_size(DEFAULT_ARENA_SIZE)
    }

    #[inline]
    pub fn with_size(arena_size: usize) -> raw_sys::Result<Self> {
        let config = PageConfig::get();
        assert!(arena_size > 0);
        assert_eq!(arena_size % config.size(), 0);

        let page_count = page_count(arena_size) as u32;
        let fd = platform::memfd_create(c"rust-mesh-alloc".as_ptr().cast(), raw_sys::MFD_CLOEXEC)?;
        platform::ftruncate(fd, arena_size as u64)?;

        let base = unsafe {
            platform::mmap(
                null_mut(),
                arena_size,
                raw_sys::PROT_READ | raw_sys::PROT_WRITE,
                raw_sys::MAP_SHARED,
                fd,
                0,
            )?
        };

        let owner_bytes = page_count as usize * size_of::<AtomicU32>();
        let owners = unsafe {
            platform::map_anonymous(owner_bytes, raw_sys::PROT_READ | raw_sys::PROT_WRITE)?
                as *mut AtomicU32
        };
        let free_span_bytes = MAX_FREE_SPANS * size_of::<Span>();
        let free_spans = unsafe {
            platform::map_anonymous(free_span_bytes, raw_sys::PROT_READ | raw_sys::PROT_WRITE)?
                as *mut Span
        };

        Ok(Self {
            config,
            arena_size,
            page_count,
            fd,
            base,
            owners,
            next_page: AtomicU32::new(0),
            free_spans,
            free_span_count: UnsafeCell::new(0),
            free_span_lock: FutexMutex::new(),
        })
    }

    #[inline(always)]
    pub const fn config(&self) -> PageConfig {
        self.config
    }

    #[inline(always)]
    pub const fn arena_size(&self) -> usize {
        self.arena_size
    }

    #[inline(always)]
    pub const fn base_ptr(&self) -> *mut u8 {
        self.base
    }

    #[inline(always)]
    pub fn contains(&self, ptr: *const u8) -> bool {
        let start = self.base as usize;
        let end = start + self.arena_size;
        let value = ptr as usize;
        start <= value && value < end
    }

    #[inline]
    pub fn reserve_pages(&self, page_count: u32, page_alignment: u32) -> Option<Span> {
        assert!(page_count > 0);
        assert!(page_alignment > 0);
        let alignment = page_alignment.next_power_of_two();

        if let Some(span) = self.take_free_span(page_count, alignment) {
            return Some(span);
        }

        loop {
            let current = self.next_page.load(Ordering::Acquire);
            let aligned = align_up_u32(current, alignment);
            let end = aligned.checked_add(page_count)?;
            if end > self.page_count {
                return None;
            }

            match self
                .next_page
                .compare_exchange(current, end, Ordering::AcqRel, Ordering::Acquire)
            {
                Ok(_) => return Some(Span::new(aligned, page_count)),
                Err(_) => continue,
            }
        }
    }

    #[inline]
    pub fn page_alloc(&self, page_count: u32, page_alignment: u32) -> Option<(*mut u8, Span)> {
        let span = self.reserve_pages(page_count, page_alignment)?;
        Some((self.ptr_from_offset(span.offset as usize), span))
    }

    #[inline]
    pub fn allocate_bytes(&self, size: usize, page_alignment: u32) -> Option<(*mut u8, Span)> {
        let pages = page_count(size) as u32;
        self.page_alloc(pages, page_alignment)
    }

    #[inline(always)]
    pub fn reserved_pages(&self) -> u32 {
        self.next_page.load(Ordering::Acquire)
    }

    pub fn reusable_span_stats(&self) -> (u32, u32) {
        let _guard = self.free_span_lock.lock();
        let count = unsafe { *self.free_span_count.get() };
        let mut pages = 0u32;
        let mut index = 0usize;
        while index < count as usize {
            let span = unsafe { *self.free_spans.add(index) };
            pages += span.length;
            index += 1;
        }
        (count, pages)
    }

    #[inline]
    pub fn track_miniheap(&self, span: Span, id: MiniHeapId) {
        for page in 0..span.length {
            self.owner_at_offset(span.offset + page)
                .store(id.value(), Ordering::Release);
        }
    }

    #[inline]
    pub fn clear_miniheap(&self, span: Span) {
        for page in 0..span.length {
            self.owner_at_offset(span.offset + page)
                .store(0, Ordering::Release);
        }
    }

    #[inline]
    pub fn release_span(&self, span: Span) {
        if span.empty() {
            return;
        }

        let _guard = self.free_span_lock.lock();
        let count = unsafe { &mut *self.free_span_count.get() };
        let mut merged = span;
        let mut index = 0usize;
        while index < *count as usize {
            let other = unsafe { *self.free_spans.add(index) };
            if other.offset + other.length == merged.offset {
                merged = Span::new(other.offset, other.length + merged.length);
                self.remove_free_span_at(index, count);
                continue;
            }
            if merged.offset + merged.length == other.offset {
                merged = Span::new(merged.offset, merged.length + other.length);
                self.remove_free_span_at(index, count);
                continue;
            }
            index += 1;
        }

        self.push_free_span(merged, count);
    }

    #[inline]
    pub fn miniheap_id_for_ptr(&self, ptr: *const u8) -> Option<MiniHeapId> {
        if !self.contains(ptr) {
            return None;
        }
        let off = self.offset_for(ptr);
        let value = self.owner_at_offset(off).load(Ordering::Acquire);
        if value == 0 {
            None
        } else {
            Some(MiniHeapId::new(value))
        }
    }

    /// # Safety
    ///
    /// `remove..remove+size` must describe a valid, page-aligned mapping within this arena.
    /// The returned alias is a private scratch mapping of the old source backing. The original
    /// `remove` range is protected with `PROT_NONE` and must be restored or remapped by the
    /// caller before any blocked mutators are allowed to resume.
    #[inline]
    pub unsafe fn begin_mesh(&self, remove: *mut u8, size: usize) -> raw_sys::Result<*mut u8> {
        let rounded = round_up_to_page(size);
        let remove_off = self.offset_for(remove);
        unsafe {
            platform::mprotect(remove, rounded, raw_sys::PROT_NONE)?;
            platform::mmap(
                core::ptr::null_mut(),
                rounded,
                raw_sys::PROT_READ,
                raw_sys::MAP_SHARED,
                self.fd,
                (remove_off as usize * self.config.size()) as u64,
            )
        }
    }

    /// # Safety
    ///
    /// Restores the source mapping to its original backing after a failed mesh attempt.
    #[inline]
    pub unsafe fn abort_mesh(
        &self,
        remove: *mut u8,
        scratch: *mut u8,
        size: usize,
    ) -> raw_sys::Result<()> {
        let rounded = round_up_to_page(size);
        let remove_off = self.offset_for(remove);
        unsafe {
            if !scratch.is_null() {
                platform::munmap(scratch, rounded)?;
            }
            platform::mmap(
                remove,
                rounded,
                raw_sys::PROT_READ | raw_sys::PROT_WRITE,
                raw_sys::MAP_SHARED | raw_sys::MAP_FIXED,
                self.fd,
                (remove_off as usize * self.config.size()) as u64,
            )?;
        }
        Ok(())
    }

    /// # Safety
    ///
    /// `keep` and `remove` must each point to valid page-aligned ranges of at least `size`
    /// bytes within this arena. The caller must ensure that aliasing these ranges is valid for
    /// the current allocator state and that any required object copying has already completed.
    #[inline]
    pub unsafe fn finalize_mesh(
        &self,
        keep: *mut u8,
        remove: *mut u8,
        scratch: *mut u8,
        size: usize,
    ) -> raw_sys::Result<()> {
        let rounded = round_up_to_page(size);
        let keep_off = self.offset_for(keep);
        let remove_off = self.offset_for(remove);
        let pages = page_count(rounded);

        unsafe {
            platform::mmap(
                remove,
                rounded,
                raw_sys::PROT_READ | raw_sys::PROT_WRITE,
                raw_sys::MAP_SHARED | raw_sys::MAP_FIXED,
                self.fd,
                (keep_off as usize * self.config.size()) as u64,
            )?;
            if !scratch.is_null() {
                platform::munmap(scratch, rounded)?;
            }
        }

        let keep_id = self.owner_at_offset(keep_off).load(Ordering::Acquire);
        for page in 0..pages {
            self.owner_at_offset(remove_off + page as u32)
                .store(keep_id, Ordering::Release);
        }

        Ok(())
    }

    #[inline]
    pub fn free_phys(&self, ptr: *mut u8, size: usize) -> raw_sys::Result<()> {
        let rounded = round_up_to_page(size);
        let offset = (ptr as usize).wrapping_sub(self.base as usize);
        platform::fallocate(
            self.fd,
            raw_sys::FALLOC_FL_PUNCH_HOLE | raw_sys::FALLOC_FL_KEEP_SIZE,
            offset as u64,
            rounded as u64,
        )
    }

    #[inline]
    pub fn reset_identity_mapping(&self, span: Span) -> raw_sys::Result<()> {
        let ptr = self.ptr_from_offset(span.offset as usize);
        unsafe {
            platform::mmap(
                ptr,
                span.byte_length_for_page_size(self.config.size()),
                raw_sys::PROT_READ | raw_sys::PROT_WRITE,
                raw_sys::MAP_SHARED | raw_sys::MAP_FIXED,
                self.fd,
                (span.offset as usize * self.config.size()) as u64,
            )?;
        }
        Ok(())
    }

    #[inline(always)]
    pub fn offset_for(&self, ptr: *const u8) -> u32 {
        let delta = (ptr as usize).wrapping_sub(self.base as usize);
        (delta >> self.config.shift()) as u32
    }

    #[inline(always)]
    pub fn ptr_from_offset(&self, offset: usize) -> *mut u8 {
        unsafe { self.base.add(offset << self.config.shift()) }
    }

    #[inline(always)]
    fn owner_at_offset(&self, offset: u32) -> &AtomicU32 {
        assert!(offset < self.page_count);
        unsafe { &*self.owners.add(offset as usize) }
    }
}

impl Drop for Arena {
    fn drop(&mut self) {
        let owner_bytes = self.page_count as usize * size_of::<AtomicU32>();
        let free_span_bytes = MAX_FREE_SPANS * size_of::<Span>();
        unsafe {
            let _ = platform::munmap(self.free_spans as *mut u8, free_span_bytes);
            let _ = platform::munmap(self.owners as *mut u8, owner_bytes);
            let _ = platform::munmap(self.base, self.arena_size);
        }
        let _ = platform::close(self.fd);
    }
}

#[inline(always)]
fn align_up_u32(value: u32, alignment: u32) -> u32 {
    debug_assert!(alignment.is_power_of_two());
    (value + alignment - 1) & !(alignment - 1)
}

impl Arena {
    fn take_free_span(&self, page_count: u32, alignment: u32) -> Option<Span> {
        let _guard = self.free_span_lock.lock();
        let count = unsafe { &mut *self.free_span_count.get() };
        let mut index = 0usize;
        while index < *count as usize {
            let span = unsafe { *self.free_spans.add(index) };
            let aligned = align_up_u32(span.offset, alignment);
            let prefix = aligned.checked_sub(span.offset)?;
            let total = prefix.checked_add(page_count)?;
            if total <= span.length {
                self.remove_free_span_at(index, count);

                if prefix > 0 {
                    self.push_free_span(Span::new(span.offset, prefix), count);
                }

                let suffix_offset = aligned + page_count;
                let suffix_length = span.length - total;
                if suffix_length > 0 {
                    self.push_free_span(Span::new(suffix_offset, suffix_length), count);
                }

                return Some(Span::new(aligned, page_count));
            }
            index += 1;
        }

        None
    }

    fn push_free_span(&self, span: Span, count: &mut u32) {
        assert!((*count as usize) < MAX_FREE_SPANS);
        unsafe {
            self.free_spans.add(*count as usize).write(span);
        }
        *count += 1;
    }

    fn remove_free_span_at(&self, index: usize, count: &mut u32) {
        debug_assert!(index < *count as usize);
        let last = *count as usize - 1;
        if index != last {
            let replacement = unsafe { *self.free_spans.add(last) };
            unsafe {
                self.free_spans.add(index).write(replacement);
            }
        }
        *count -= 1;
    }
}
