use core::cmp::Ordering;

use alloc::collections::{linked_list::CursorMut, LinkedList};
use yggdrasil_abi::error::Error;

use crate::VirtualMemoryRange;

#[derive(PartialEq, Clone, Debug, Copy)]
struct AllocatorNode {
    range: VirtualMemoryRange,
    used: bool,
}

impl PartialOrd for AllocatorNode {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.range.start_pfn.partial_cmp(&other.range.start_pfn)
    }
}

pub struct TreeAllocator {
    ranges: LinkedList<AllocatorNode>,
}

impl AllocatorNode {
    pub const fn free(start_pfn: usize, end_pfn: usize) -> Self {
        Self {
            range: VirtualMemoryRange { start_pfn, end_pfn },
            used: false,
        }
    }

    #[cfg(test)]
    pub const fn used(start_pfn: usize, end_pfn: usize) -> Self {
        Self {
            range: VirtualMemoryRange { start_pfn, end_pfn },
            used: true,
        }
    }

    #[inline]
    pub const fn pfn_count(&self) -> usize {
        self.range.end_pfn - self.range.start_pfn
    }
}

impl TreeAllocator {
    pub fn new(start_pfn: usize, end_pfn: usize) -> Self {
        let mut ranges = LinkedList::new();
        ranges.push_back(AllocatorNode::free(start_pfn, end_pfn));
        Self { ranges }
    }

    fn find_region_mut<F: Fn(&AllocatorNode) -> bool>(&mut self, f: F) -> CursorMut<AllocatorNode> {
        let mut cursor = self.ranges.cursor_front_mut();
        while let Some(range) = cursor.current() {
            if f(range) {
                break;
            }
            cursor.move_next();
        }
        cursor
    }

    fn coalesce_regions(&mut self) {
        let mut cursor = self.ranges.cursor_front_mut();

        loop {
            let Some(&mut next) = cursor.peek_next() else {
                break;
            };
            let current = cursor.current().unwrap();

            if current.used == next.used {
                debug_assert_eq!(current.range.end_pfn, next.range.start_pfn);
                current.range.end_pfn = next.range.end_pfn;

                cursor.move_next();
                cursor.remove_current();
                cursor.move_prev();
            } else {
                cursor.move_next();
            }
        }
    }

    fn set_range(
        &mut self,
        start_pfn: usize,
        pfn_count: usize,
        old_state: bool,
        new_state: bool,
    ) -> Result<(), Error> {
        let insert = VirtualMemoryRange {
            start_pfn,
            end_pfn: start_pfn + pfn_count,
        };
        let mut cursor = self.find_region_mut(|r| r.used == old_state && r.range.contains(&insert));
        let range = cursor.current().ok_or(Error::AlreadyExists)?;

        let start_pfn = range.range.start_pfn;
        let end_pfn = range.range.end_pfn;

        match (insert.start_pfn == start_pfn, insert.end_pfn == end_pfn) {
            // No split
            (true, true) => {
                range.used = new_state;
            }
            // Split start
            (true, false) => {
                range.used = new_state;
                range.range.end_pfn = insert.end_pfn;

                cursor.insert_after(AllocatorNode {
                    range: VirtualMemoryRange {
                        start_pfn: insert.end_pfn,
                        end_pfn,
                    },
                    used: old_state,
                });
            }
            // Split end
            (false, true) => {
                range.range.end_pfn = insert.start_pfn;

                cursor.insert_after(AllocatorNode {
                    range: VirtualMemoryRange {
                        start_pfn: insert.start_pfn,
                        end_pfn,
                    },
                    used: new_state,
                });
            }
            // Split in the middle
            (false, false) => {
                range.range = insert;
                range.used = new_state;

                cursor.insert_after(AllocatorNode {
                    range: VirtualMemoryRange {
                        start_pfn: insert.end_pfn,
                        end_pfn,
                    },
                    used: old_state,
                });
                cursor.insert_before(AllocatorNode {
                    range: VirtualMemoryRange {
                        start_pfn,
                        end_pfn: insert.start_pfn,
                    },
                    used: old_state,
                });
            }
        }

        self.coalesce_regions();

        Ok(())
    }

    pub fn insert(&mut self, start_pfn: usize, pfn_count: usize) -> Result<(), Error> {
        self.set_range(start_pfn, pfn_count, false, true)
    }

    pub fn free(&mut self, start_pfn: usize, pfn_count: usize) -> Result<(), Error> {
        self.set_range(start_pfn, pfn_count, true, false)
    }

    pub fn allocate(&mut self, pfn_count: usize) -> Option<usize> {
        let mut cursor = self.find_region_mut(|r| !r.used && r.pfn_count() >= pfn_count);
        let range = cursor.current()?;

        let start_pfn = range.range.start_pfn;
        let end_pfn = range.range.end_pfn;

        range.used = true;

        if range.pfn_count() > pfn_count {
            range.range.end_pfn = start_pfn + pfn_count;

            // Split the range
            cursor.insert_after(AllocatorNode::free(start_pfn + pfn_count, end_pfn));
        }

        self.coalesce_regions();

        Some(start_pfn)
    }

    pub fn ranges(&self) -> impl Iterator<Item = (bool, VirtualMemoryRange)> + '_ {
        self.ranges.iter().map(|r| (r.used, r.range))
    }
}

#[cfg(test)]
mod tests {
    use alloc::collections::LinkedList;

    use super::{AllocatorNode, TreeAllocator};

    extern crate std;

    #[test]
    fn deallocation() {
        let ranges = LinkedList::from_iter([
            AllocatorNode::free(0, 12),
            AllocatorNode::used(12, 24),
            AllocatorNode::free(24, 32),
            AllocatorNode::used(32, 64),
            AllocatorNode::free(64, 128),
        ]);
        let mut alloc = TreeAllocator { ranges };

        // No-split dealloc
        assert_eq!(alloc.free(12, 12), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::free(0, 32),
            AllocatorNode::used(32, 64),
            AllocatorNode::free(64, 128),
        ]);
        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Split at the start dealloc
        assert_eq!(alloc.free(32, 8), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::free(0, 40),
            AllocatorNode::used(40, 64),
            AllocatorNode::free(64, 128),
        ]);
        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Split at the end dealloc
        assert_eq!(alloc.free(56, 8), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::free(0, 40),
            AllocatorNode::used(40, 56),
            AllocatorNode::free(56, 128),
        ]);

        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Split in the middle
        assert_eq!(alloc.free(42, 4), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::free(0, 40),
            AllocatorNode::used(40, 42),
            AllocatorNode::free(42, 46),
            AllocatorNode::used(46, 56),
            AllocatorNode::free(56, 128),
        ]);

        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Whole region free
        assert_eq!(alloc.free(40, 2), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::free(0, 46),
            AllocatorNode::used(46, 56),
            AllocatorNode::free(56, 128),
        ]);

        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        assert_eq!(alloc.free(46, 10), Ok(()));
        let expected = LinkedList::from_iter([AllocatorNode::free(0, 128)]);

        itertools::assert_equal(alloc.ranges.iter(), expected.iter());
    }

    #[test]
    fn allocation() {
        let ranges = LinkedList::from_iter([
            AllocatorNode::free(0, 12),
            AllocatorNode::used(12, 24),
            AllocatorNode::free(24, 32),
            AllocatorNode::used(32, 64),
        ]);
        let mut alloc = TreeAllocator { ranges };

        // Non-splitting allocation
        assert_eq!(alloc.allocate(12), Some(0));

        // Splitting allocation
        assert_eq!(alloc.allocate(4), Some(24));

        // Non-splitting allocation
        assert_eq!(alloc.allocate(4), Some(28));

        // Out of memory
        assert_eq!(alloc.allocate(1), None);

        let expected = LinkedList::from_iter([AllocatorNode::used(0, 64)]);

        itertools::assert_equal(alloc.ranges, expected);
    }

    #[test]
    fn insertion() {
        let ranges = LinkedList::from_iter([
            AllocatorNode::free(0, 12),
            AllocatorNode::used(12, 24),
            AllocatorNode::free(24, 32),
            AllocatorNode::used(32, 64),
            AllocatorNode::free(64, 128),
        ]);
        let mut alloc = TreeAllocator { ranges };

        // No split
        assert_eq!(alloc.insert(0, 12), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::used(0, 24),
            AllocatorNode::free(24, 32),
            AllocatorNode::used(32, 64),
            AllocatorNode::free(64, 128),
        ]);
        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Split at the start
        assert_eq!(alloc.insert(24, 4), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::used(0, 28),
            AllocatorNode::free(28, 32),
            AllocatorNode::used(32, 64),
            AllocatorNode::free(64, 128),
        ]);
        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Split at the end
        assert_eq!(alloc.insert(30, 2), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::used(0, 28),
            AllocatorNode::free(28, 30),
            AllocatorNode::used(30, 64),
            AllocatorNode::free(64, 128),
        ]);
        itertools::assert_equal(alloc.ranges.iter(), expected.iter());

        // Split in the middle
        assert_eq!(alloc.insert(72, 16), Ok(()));
        let expected = LinkedList::from_iter([
            AllocatorNode::used(0, 28),
            AllocatorNode::free(28, 30),
            AllocatorNode::used(30, 64),
            AllocatorNode::free(64, 72),
            AllocatorNode::used(72, 88),
            AllocatorNode::free(88, 128),
        ]);
        itertools::assert_equal(alloc.ranges.iter(), expected.iter());
    }
}
