use embassy_time::Timer;
use esp_alloc::{HeapRegion, MemoryCapability};
use esp_hal::ram;
use log::{info, warn};

// Memory allocation regions.
// These can be debugged using `xtensa-esp32s3-elf-size -A <path-to-binary>`.
// A panic such as `memory allocation of 3740121773 bytes failed` is caused by a heap overflow. The size is `DEEDBAAD` in hex.
//
// RAM usage of static variables can be diagnosed with:
// ```sh
// xtensa-esp32s3-elf-nm.exe -S --size-sort -t d ../target/xtensa-esp32s3-none-elf/release/acid-firmware | grep -iE ' [dbv] ' | tail -n 10 | tac
// ```

/// Total heap size
pub const HEAP_SIZE: usize = 196 * 1024;
/// Size of the app core's stack
pub const STACK_SIZE_CORE_APP: usize = 80 * 1024;

pub static PSRAM_ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();
pub type PsramAllocator = &'static esp_alloc::EspHeap;

pub fn initialize(psram_peripheral: esp_hal::peripherals::PSRAM) {
    // Use the internal DRAM as the heap.
    // Memory reclaimed from the esp-idf bootloader.
    const HEAP_SIZE_RECLAIMED: usize = const {
        let range = esp_metadata_generated::memory_range!("DRAM2_UNINIT");
        range.end - range.start
    };

    esp_alloc::heap_allocator!(#[ram(reclaimed)] size: HEAP_SIZE_RECLAIMED);
    esp_alloc::heap_allocator!(size: HEAP_SIZE - HEAP_SIZE_RECLAIMED);

    #[cfg(not(feature = "no-alloc-tracing"))]
    alloc_tracing::install_ram_allocator_proxy();

    info!(
        "Internal RAM heap initialized!\n{}",
        esp_alloc::HEAP.stats()
    );

    // Initialize the PSRAM allocator.
    {
        let (psram_offset, psram_size) = esp_hal::psram::psram_raw_parts(&psram_peripheral);
        unsafe {
            PSRAM_ALLOCATOR.add_region(HeapRegion::new(
                psram_offset,
                psram_size,
                MemoryCapability::External.into(),
            ));
        }
        info!(
            "External PSRAM heap initialized with capacity of {} MiB!\n{}",
            psram_size / 1024 / 1024,
            PSRAM_ALLOCATOR.stats(),
        );
    }
}

pub async fn run_alloc_stats_reporter() {
    let mut psram_used_prev = 0;
    let mut heap_used_prev = 0;
    loop {
        let psram_stats = PSRAM_ALLOCATOR.stats();
        let heap_stats = esp_alloc::HEAP.stats();
        if psram_stats.current_usage != psram_used_prev {
            let difference = psram_stats.current_usage as isize - psram_used_prev as isize;
            psram_used_prev = psram_stats.current_usage;
            warn!(
                "External PSRAM heap usage changed: {}{}\n{psram_stats}",
                if difference < 0 { '-' } else { '+' },
                difference.abs()
            );
        }
        if heap_stats.current_usage != heap_used_prev {
            let difference = heap_stats.current_usage as isize - heap_used_prev as isize;
            heap_used_prev = heap_stats.current_usage;
            warn!(
                "Internal RAM heap usage changed: {}{}\n{heap_stats}",
                if difference < 0 { '-' } else { '+' },
                difference.abs()
            );
            #[cfg(not(feature = "no-alloc-tracing"))]
            alloc_tracing::report_stats();
        }
        Timer::after_secs(1).await;
    }
}

#[cfg(not(feature = "no-alloc-tracing"))]
mod alloc_tracing {
    use core::{
        alloc::GlobalAlloc,
        cell::RefCell,
        cmp::{Ordering, Reverse},
        fmt::Display,
    };

    use alloc::collections::btree_map::BTreeMap;
    use embassy_sync::blocking_mutex::{Mutex, raw::CriticalSectionRawMutex};
    use esp_alloc::EspHeap;
    use esp_backtrace::Backtrace;
    use esp_sync::NonReentrantMutex;
    use log::warn;
    use tinyvec::ArrayVec;

    use crate::ram::{PSRAM_ALLOCATOR, PsramAllocator};

    #[derive(Default, Clone, PartialEq, Eq, Debug)]
    struct Stats {
        /// The total number of allocations.
        allocations: usize,
        /// The number of bytes allocated in total.
        allocated_total: usize,
        /// The number of bytes allocated minus the number of bytes deallocated.
        allocated_current: usize,
    }

    impl Stats {
        fn update(&mut self, allocations_delta: isize, bytes_delta: isize, bytes_new: usize) {
            self.allocations = (self.allocations as isize + allocations_delta) as usize;
            self.allocated_total += bytes_new;
            self.allocated_current = (self.allocated_current as isize + bytes_delta) as usize;
        }
    }

    impl PartialOrd for Stats {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
            Some(self.cmp(other))
        }
    }

    impl Ord for Stats {
        fn cmp(&self, other: &Self) -> Ordering {
            self.allocated_current
                .cmp(&other.allocated_current)
                .then_with(|| self.allocated_total.cmp(&other.allocated_total).reverse())
                .then_with(|| self.allocations.cmp(&other.allocations).reverse())
        }
    }

    #[derive(Clone)]
    struct BacktraceWrapper(Backtrace);

    impl Display for BacktraceWrapper {
        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
            let mut it = self.0.frames().iter();
            if let Some(frame) = it.next() {
                write!(f, "0x{:08x}", frame.program_counter())?;
            }
            while let Some(frame) = it.next() {
                write!(f, " 0x{:08x}", frame.program_counter())?;
            }
            Ok(())
        }
    }

    impl PartialEq for BacktraceWrapper {
        fn eq(&self, other: &Self) -> bool {
            if self.0.frames().len() != other.0.frames().len() {
                return false;
            }

            for (lhs_frame, rhs_frame) in self.0.frames().iter().zip(other.0.frames()) {
                if lhs_frame.program_counter() != rhs_frame.program_counter() {
                    return false;
                }
            }

            true
        }
    }

    impl Eq for BacktraceWrapper {}

    impl PartialOrd for BacktraceWrapper {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
            Some(self.cmp(other))
        }
    }

    impl Ord for BacktraceWrapper {
        fn cmp(&self, other: &Self) -> Ordering {
            let mut lhs_it = self.0.frames().iter().rev();
            let mut rhs_it = other.0.frames().iter().rev();

            loop {
                match (lhs_it.next(), rhs_it.next()) {
                    (Some(lhs_frame), Some(rhs_frame)) => {
                        let ordering = lhs_frame
                            .program_counter()
                            .cmp(&rhs_frame.program_counter());

                        if ordering != Ordering::Equal {
                            return ordering;
                        }
                    }
                    (Some(_), None) => return Ordering::Greater,
                    (None, Some(_)) => return Ordering::Less,
                    (None, None) => return Ordering::Equal,
                }
            }
        }
    }

    #[derive(Clone)]
    struct AllocationTracer {
        bt_to_stats: BTreeMap<BacktraceWrapper, Stats, PsramAllocator>,
        ptr_to_bt: BTreeMap<*mut u8, BacktraceWrapper, PsramAllocator>,
    }

    unsafe impl Send for AllocationTracer {}

    impl AllocationTracer {
        const fn new() -> Self {
            Self {
                bt_to_stats: BTreeMap::new_in(&PSRAM_ALLOCATOR),
                ptr_to_bt: BTreeMap::new_in(&PSRAM_ALLOCATOR),
            }
        }

        fn alloc(&mut self, bt: Backtrace, ptr: *mut u8, bytes: usize) {
            let bt = BacktraceWrapper(bt.clone());
            self.ptr_to_bt.insert(ptr, bt.clone());
            let stats = self
                .bt_to_stats
                .entry(bt)
                .or_insert_with(|| Default::default());
            stats.update(1, bytes as isize, bytes);
        }

        fn realloc(
            &mut self,
            ptr_old: *mut u8,
            ptr_new: *mut u8,
            bytes_old: usize,
            bytes_new: usize,
        ) {
            let bt = self.ptr_to_bt.remove(&ptr_old).unwrap();
            self.ptr_to_bt.insert(ptr_new, bt.clone());
            let stats = self.bt_to_stats.get_mut(&bt).unwrap();
            stats.update(0, bytes_new as isize - bytes_old as isize, bytes_new);
        }

        fn dealloc(&mut self, ptr: *mut u8, bytes: usize) {
            let bt = self.ptr_to_bt.remove(&ptr).unwrap();
            let stats = self.bt_to_stats.get_mut(&bt).unwrap();
            stats.update(-1, -(bytes as isize), 0);
        }
    }

    struct TracingAllocator<T: GlobalAlloc + 'static> {
        inner: NonReentrantMutex<Option<&'static T>>,
        tracer: Mutex<CriticalSectionRawMutex, RefCell<AllocationTracer>>,
    }

    impl<T> TracingAllocator<T>
    where
        T: GlobalAlloc,
    {
        const fn new() -> Self {
            Self {
                inner: NonReentrantMutex::new(None),
                tracer: Mutex::new(RefCell::new(AllocationTracer::new())),
            }
        }

        fn with_inner<R>(&self, callback: impl FnOnce(&T) -> R) -> R {
            self.inner.with(|inner| {
                (callback)(
                    inner
                        .as_ref()
                        .expect("an allocator must be installed in the global allocator proxy"),
                )
            })
        }

        fn update_with<R>(&self, callback: impl FnOnce(&mut AllocationTracer) -> R) -> R {
            self.tracer
                .lock(|tracer| (callback)(&mut tracer.borrow_mut()))
        }
    }

    unsafe impl<T> GlobalAlloc for TracingAllocator<T>
    where
        T: GlobalAlloc,
    {
        unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
            let bt = Backtrace::capture();
            let ptr = self.with_inner(|inner| unsafe { inner.alloc(layout) });
            self.update_with(|tracer| tracer.alloc(bt, ptr, layout.size()));
            ptr
        }

        unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 {
            let bt = Backtrace::capture();
            let ptr = self.with_inner(|inner| unsafe { inner.alloc_zeroed(layout) });
            self.update_with(|tracer| tracer.alloc(bt, ptr, layout.size()));
            ptr
        }

        unsafe fn realloc(
            &self,
            ptr_old: *mut u8,
            layout: core::alloc::Layout,
            new_size: usize,
        ) -> *mut u8 {
            let ptr_new =
                self.with_inner(|inner| unsafe { inner.realloc(ptr_old, layout, new_size) });
            self.update_with(|tracer| tracer.realloc(ptr_old, ptr_new, layout.size(), new_size));
            ptr_new
        }

        unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) {
            self.update_with(|tracer| tracer.dealloc(ptr, layout.size()));
            self.with_inner(|inner| unsafe { inner.dealloc(ptr, layout) });
        }
    }

    #[global_allocator]
    static PROXY: TracingAllocator<EspHeap> = TracingAllocator::new();

    pub fn install_ram_allocator_proxy() {
        PROXY.inner.with(|inner| {
            *inner = Some(&esp_alloc::HEAP);
        });
    }

    pub fn report_stats() {
        let AllocationTracer {
            bt_to_stats,
            ptr_to_bt,
        } = PROXY.tracer.lock(|tracer| tracer.borrow().clone());
        // Wrapped in `Option` because of the `Default` requirement.
        let mut sorted = ArrayVec::<[Option<(BacktraceWrapper, Reverse<Stats>)>; 5]>::new();
        let bt_to_stats_len = bt_to_stats.len();

        for (key, value) in bt_to_stats {
            // Reverse ordering of stats because we are interested in the largest.
            let value_rev = Reverse(value);
            if let Some((_, value_last)) = sorted.last().and_then(|last| last.as_ref())
                && &value_rev >= value_last
                && sorted.len() >= sorted.capacity()
            {
                // This stat is not large enough to be inserted.
                continue;
            }
            let index = match sorted.binary_search_by_key(&Some(&value_rev), |item| {
                item.as_ref().map(|(_, current_value)| current_value)
            }) {
                Ok(index) => index,
                Err(index) => index,
            };
            if sorted.len() >= sorted.capacity() {
                assert!(
                    index < sorted.len(),
                    "the stat should be large enough to be inserted not at the end of the list"
                );
                let _ = sorted.pop().unwrap();
            }
            sorted.insert(index, Some((key, value_rev)));
        }

        warn!("Largest allocations in global allocator:");

        for (index, (key, value)) in sorted.into_iter().map(Option::unwrap).enumerate() {
            warn!("{}. {}\n{:#?}", index + 1, key, value.0);
        }

        warn!("bt_to_stats.len() = {}", bt_to_stats_len);
        warn!("ptr_to_bt.len() = {}", ptr_to_bt.len());
    }
}
