//! Block vector management structures
use core::{
    cmp::Ordering,
    mem::MaybeUninit,
    ops::{Index, IndexMut},
};

use yggdrasil_abi::error::Error;

use crate::block::{self, BlockAllocator, BlockData, BlockIndirect, BlockIndirectIndirect};

const L0_BLOCKS: usize = 32; // 128K in L0
const L1_BLOCKS: usize = 16; // 32M in L1

/// Block vector for efficient in-memory files
pub struct BVec<'a, A: BlockAllocator> {
    capacity: usize,
    size: usize,
    l0: [BlockData<'a, A>; L0_BLOCKS],
    l1: [BlockIndirect<'a, A>; L1_BLOCKS],
    l2: BlockIndirectIndirect<'a, A>,
}

impl<'a, A: BlockAllocator> BVec<'a, A> {
    /// Creates an empty block vector.
    ///
    /// # Note
    ///
    /// The function is guaranteed to make no allocations before the vector is actually written to.
    pub fn new() -> Self {
        let mut l0 = [const { MaybeUninit::uninit() }; L0_BLOCKS];
        let mut l1 = [const { MaybeUninit::uninit() }; L1_BLOCKS];

        for it in l0.iter_mut() {
            it.write(BlockData::null());
        }

        for it in l1.iter_mut() {
            it.write(BlockIndirect::null());
        }

        Self {
            capacity: 0,
            size: 0,
            l0: unsafe { MaybeUninit::array_assume_init(l0) },
            l1: unsafe { MaybeUninit::array_assume_init(l1) },
            l2: BlockIndirectIndirect::null(),
        }
    }

    /// Initializes the block vector with existing data, marking all blocks as Copy-on-Write
    pub fn init_with_cow(&mut self, data: &'static [u8]) -> Result<(), Error> {
        let data_ptr = data.as_ptr() as usize;
        assert_eq!(data_ptr & 1, 0);

        let blocks = data.len().div_ceil(block::SIZE);
        self.resize(blocks)?;

        for i in 0..blocks {
            let src = data_ptr + i * block::SIZE;
            let block = unsafe { BlockData::copy_on_write(src) };
            self[i] = block;
        }

        self.size = data.len();
        Ok(())
    }

    /// Returns the size of the data inside this vector
    #[inline]
    pub const fn size(&self) -> usize {
        self.size
    }

    #[inline]
    fn l1_indirect_count(cap: usize) -> usize {
        core::cmp::min(
            cap.saturating_sub(L0_BLOCKS).div_ceil(block::ENTRY_COUNT),
            L1_BLOCKS,
        )
    }

    #[inline]
    fn l2_indirect_count(cap: usize) -> usize {
        cap.saturating_sub(L0_BLOCKS + L1_BLOCKS * block::ENTRY_COUNT)
            .div_ceil(block::ENTRY_COUNT)
    }

    fn grow_l1(&mut self, old: usize, new: usize) -> Result<(), Error> {
        let old_l1_cap = Self::l1_indirect_count(old);
        let new_l1_cap = Self::l1_indirect_count(new);

        debug_assert!(old_l1_cap <= new_l1_cap);

        for l1 in &mut self.l1[old_l1_cap..new_l1_cap] {
            debug_assert!(l1.is_null());
            *l1 = BlockIndirect::new()?;
        }

        Ok(())
    }

    fn grow_l2(&mut self, old: usize, new: usize) -> Result<(), Error> {
        let old_l2_cap = Self::l2_indirect_count(old);
        let new_l2_cap = Self::l2_indirect_count(new);

        if old_l2_cap == 0 && new_l2_cap != 0 {
            // Allocate the indirect->indirect block
            self.l2 = BlockIndirectIndirect::new()?;
        }

        if old_l2_cap < new_l2_cap {
            debug_assert!(!self.l2.is_null());
            for l2 in &mut self.l2[old_l2_cap..new_l2_cap] {
                *l2 = BlockIndirect::new()?;
            }
        }

        Ok(())
    }

    fn shrink_l1(&mut self, old: usize, new: usize) {
        let old_l1_cap = Self::l1_indirect_count(old);
        let new_l1_cap = Self::l1_indirect_count(new);

        debug_assert!(old_l1_cap >= new_l1_cap);

        for l1 in &mut self.l1[new_l1_cap..old_l1_cap] {
            debug_assert!(!l1.is_null());
            *l1 = BlockIndirect::null();
        }
    }

    fn shrink_l2(&mut self, old: usize, new: usize) {
        let old_l2_cap = Self::l2_indirect_count(old);
        let new_l2_cap = Self::l2_indirect_count(new);

        debug_assert!(old_l2_cap >= new_l2_cap);

        if new_l2_cap < old_l2_cap {
            debug_assert!(!self.l2.is_null());
            for l2 in &mut self.l2[new_l2_cap..old_l2_cap] {
                l2.set_null();
            }
        }

        if new_l2_cap == 0 && old_l2_cap != 0 {
            // Free the L2 block itself
            self.l2 = BlockIndirectIndirect::null();
        }
    }

    /// Resizes the vector to hold exactly `new_capacity` data blocks
    pub fn resize(&mut self, new_capacity: usize) -> Result<(), Error> {
        if new_capacity
            > L0_BLOCKS + L1_BLOCKS * block::ENTRY_COUNT + block::ENTRY_COUNT * block::ENTRY_COUNT
        {
            log::warn!("Requested file capacity too large: {new_capacity}");
            return Err(Error::OutOfMemory);
        }

        match new_capacity.cmp(&self.capacity) {
            Ordering::Less => {
                // Shrink data blocks
                for index in new_capacity..self.capacity {
                    let block = &mut self[index];
                    debug_assert!(!block.is_null());
                    block.set_null();
                }

                // Shrink L1 and L2 blocks
                self.shrink_l1(self.capacity, new_capacity);
                self.shrink_l2(self.capacity, new_capacity);
            }

            Ordering::Greater => {
                // Allocate more L1 and L2 blocks if necessary
                self.grow_l1(self.capacity, new_capacity)?;
                self.grow_l2(self.capacity, new_capacity)?;

                // Grow data blocks
                for index in self.capacity..new_capacity {
                    let block = unsafe { self.index_unchecked_mut(index) };
                    assert!(block.is_null());
                    *block = BlockData::new()?;
                    block.fill(0);
                }
            }

            Ordering::Equal => (),
        }

        self.capacity = new_capacity;
        Ok(())
    }

    fn ensure_write_capacity(&mut self, pos: usize, need_to_write: usize) -> Result<(), Error> {
        let current_capacity = self.capacity;
        let need_capacity = core::cmp::max(pos + need_to_write, self.size).div_ceil(block::SIZE);

        if need_capacity > current_capacity {
            self.resize(need_capacity)
        } else {
            Ok(())
        }
    }

    /// Writes data to the vector, growing it if needed
    pub fn write(&mut self, pos: u64, data: &[u8]) -> Result<usize, Error> {
        let mut pos = pos as usize;
        let mut rem = data.len();
        let mut doff = 0usize;

        self.ensure_write_capacity(pos, rem)?;

        if pos + rem > self.size {
            self.size = pos + rem;
        }

        while rem > 0 {
            let index = pos / block::SIZE;
            let offset = pos % block::SIZE;
            let count = core::cmp::min(rem, block::SIZE - offset);

            let block = &mut self[index];

            let dst = &mut block[offset..offset + count];
            let src = &data[doff..doff + count];
            dst.copy_from_slice(src);

            doff += count;
            pos += count;
            rem -= count;
        }

        Ok(doff)
    }

    /// Reads data from the vector
    pub fn read(&self, pos: u64, data: &mut [u8]) -> Result<usize, Error> {
        let mut pos = pos as usize;
        if pos > self.size {
            return Err(Error::InvalidFile);
        }

        let mut rem = core::cmp::min(self.size - pos, data.len());
        let mut doff = 0usize;

        while rem > 0 {
            let index = pos / block::SIZE;
            let offset = pos % block::SIZE;
            let count = core::cmp::min(block::SIZE - offset, rem);

            let block = &self[index];

            let src = &block[offset..offset + count];
            let dst = &mut data[doff..doff + count];

            dst.copy_from_slice(src);

            doff += count;
            pos += count;
            rem -= count;
        }

        Ok(doff)
    }

    /// Resize the block vector to requested size
    pub fn truncate(&mut self, new_size: u64) -> Result<(), Error> {
        let new_size: usize = new_size.try_into().unwrap();
        let requested_capacity = new_size.div_ceil(block::SIZE);
        self.resize(requested_capacity)?;
        // TODO fill with zeros if resizing larger?
        self.size = new_size;
        Ok(())
    }

    unsafe fn index_unchecked(&self, mut index: usize) -> &BlockData<'a, A> {
        if index < L0_BLOCKS {
            return &self.l0[index];
        }
        index -= L0_BLOCKS;
        if index < L1_BLOCKS * block::ENTRY_COUNT {
            let l1i = index / block::ENTRY_COUNT;
            let l0i = index % block::ENTRY_COUNT;

            let l1r = &self.l1[l1i];
            assert!(!l1r.is_null());

            return &l1r[l0i];
        }
        index -= L1_BLOCKS * block::ENTRY_COUNT;
        if index < block::ENTRY_COUNT * block::ENTRY_COUNT {
            let l2_outer = index / block::ENTRY_COUNT;
            let l2_inner = index % block::ENTRY_COUNT;

            let l2 = &self.l2[l2_outer];
            debug_assert!(!l2.is_null());

            &l2[l2_inner]
        } else {
            panic!("Block index too large");
        }
    }

    unsafe fn index_unchecked_mut(&mut self, mut index: usize) -> &mut BlockData<'a, A> {
        if index < L0_BLOCKS {
            return &mut self.l0[index];
        }
        index -= L0_BLOCKS;
        if index < L1_BLOCKS * block::ENTRY_COUNT {
            let l1i = index / block::ENTRY_COUNT;
            let l0i = index % block::ENTRY_COUNT;

            let l1r = &mut self.l1[l1i];
            assert!(!l1r.is_null());

            return &mut l1r[l0i];
        }
        index -= L1_BLOCKS * block::ENTRY_COUNT;
        if index < block::ENTRY_COUNT * block::ENTRY_COUNT {
            let l2_outer = index / block::ENTRY_COUNT;
            let l2_inner = index % block::ENTRY_COUNT;

            let l2 = &mut self.l2[l2_outer];
            debug_assert!(!l2.is_null());

            &mut l2[l2_inner]
        } else {
            panic!("Block index too large");
        }
    }
}

impl<'a, A: BlockAllocator> Index<usize> for BVec<'a, A> {
    type Output = BlockData<'a, A>;

    fn index(&self, index: usize) -> &Self::Output {
        if index > self.capacity {
            panic!(
                "Block index out of bounds: capacity={}, index={}",
                self.capacity, index
            );
        }

        unsafe { self.index_unchecked(index) }
    }
}

impl<A: BlockAllocator> IndexMut<usize> for BVec<'_, A> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        if index > self.capacity {
            panic!(
                "Block index out of bounds: capacity={}, index={}",
                self.capacity, index
            );
        }

        unsafe { self.index_unchecked_mut(index) }
    }
}

impl<A: BlockAllocator> TryFrom<&'static [u8]> for BVec<'_, A> {
    type Error = Error;

    fn try_from(value: &'static [u8]) -> Result<Self, Self::Error> {
        let mut res = Self::new();
        res.init_with_cow(value)?;
        assert_eq!(res.size(), value.len());
        Ok(res)
    }
}

#[cfg(test)]
mod bvec_allocation {
    use core::sync::atomic::Ordering;

    use crate::{
        block,
        bvec::{BVec, L0_BLOCKS, L1_BLOCKS},
    };

    #[test]
    fn bvec_test_capacity() {
        test_allocator_with_counter!(A_COUNTER, A);

        let capacity = 200000; // ~781MiB file
        assert_eq!(BVec::<A>::l1_indirect_count(capacity), L1_BLOCKS);
        assert_eq!(BVec::<A>::l2_indirect_count(capacity), 375);
    }

    #[test]
    fn bvec_grow_shrink_l2() {
        test_allocator_with_counter!(A_COUNTER, A);

        assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);

        {
            let mut bvec = BVec::<A>::new();

            // L0 + L1 + 1000 L2 blocks
            const N: usize = L0_BLOCKS + L1_BLOCKS * block::ENTRY_COUNT + 1000;

            bvec.resize(N).unwrap();
            assert_eq!(A_COUNTER.load(Ordering::Acquire), N + L1_BLOCKS + 1 + 2);

            // Grow some more
            const M: usize = N + 10000;

            bvec.resize(M).unwrap();
            assert_eq!(A_COUNTER.load(Ordering::Acquire), M + L1_BLOCKS + 1 + 22);

            // Shrink back to N
            bvec.resize(N).unwrap();
            assert_eq!(A_COUNTER.load(Ordering::Acquire), N + L1_BLOCKS + 1 + 2);

            // Shrink to less than L2
            const O: usize = L0_BLOCKS + 3;

            bvec.resize(O).unwrap();
            assert_eq!(A_COUNTER.load(Ordering::Acquire), O + 1);
        }
    }

    #[test]
    fn bvec_grow_shrink() {
        test_allocator_with_counter!(A_COUNTER, A);

        assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);

        {
            let mut bvec = BVec::<A>::new();

            assert_eq!(
                A_COUNTER.load(Ordering::Acquire),
                0,
                "BVec should not allocate on creation"
            );

            const N: usize = 123;
            bvec.resize(N).unwrap();

            // N data blocks (12 in L0 + 111 in L1)
            assert_eq!(A_COUNTER.load(Ordering::Acquire), N + 1);

            // Test the index interface
            for i in 0..N {
                assert!(!bvec[i].is_null(), "Index {} must be allocated", i);
            }

            // Test the data structure
            for i in 0..L0_BLOCKS {
                assert!(!bvec.l0[i].is_null());
            }
            assert!(!bvec.l1[0].is_null());
            for i in L0_BLOCKS..N {
                let l1i = (i - L0_BLOCKS) / block::ENTRY_COUNT;
                let l0i = (i - L0_BLOCKS) % block::ENTRY_COUNT;

                let l1r = &bvec.l1[l1i];
                assert!(!l1r.is_null());
                assert!(!l1r[l0i].is_null());
            }

            for i in 1..L1_BLOCKS {
                assert!(bvec.l1[i].is_null());
            }

            // Shrink to 100 blocks, test if L1 is still allocated
            const M: usize = 100;
            bvec.resize(M).unwrap();

            assert_eq!(A_COUNTER.load(Ordering::Acquire), M + 1);

            // Test the index interface
            for i in 0..M {
                assert!(!bvec[i].is_null(), "Index {} must be allocated", i);
            }

            // Test the data structure
            for i in 0..L0_BLOCKS {
                assert!(!bvec.l0[i].is_null());
            }
            assert!(!bvec.l1[0].is_null());
            for i in L0_BLOCKS..M {
                let l1i = (i - L0_BLOCKS) / block::ENTRY_COUNT;
                let l0i = (i - L0_BLOCKS) % block::ENTRY_COUNT;

                let l1r = &bvec.l1[l1i];
                assert!(!l1r.is_null());
                assert!(!l1r[l0i].is_null());
            }
            for i in M..N {
                let l1i = (i - L0_BLOCKS) / block::ENTRY_COUNT;
                let l0i = (i - L0_BLOCKS) % block::ENTRY_COUNT;

                let l1r = &bvec.l1[l1i];
                assert!(!l1r.is_null());
                assert!(l1r[l0i].is_null());
            }

            for i in 1..L1_BLOCKS {
                assert!(bvec.l1[i].is_null());
            }

            // Shrink to 13 blocks, test if L1 got deallocated
            const O: usize = 13;
            bvec.resize(O).unwrap();

            assert_eq!(A_COUNTER.load(Ordering::Acquire), O);
        }

        assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);
    }
}

#[cfg(all(test, any(feature = "test-io", rust_analyzer)))]
mod bvec_io {
    use crate::{block, bvec::L0_BLOCKS};

    use super::BVec;

    #[test]
    fn test_bvec_write() {
        test_allocator_with_counter!(A_COUNTER, A);

        {
            let data = [1, 2, 3, 4, 5];
            let mut bvec = BVec::<A>::new();

            // Write at 0
            assert_eq!(bvec.write(0, &data).unwrap(), data.len());
            assert_eq!(bvec.capacity, 1);
            assert_eq!(bvec.size(), data.len());

            assert_eq!(&bvec[0][..bvec.size()], &data[..]);

            // Write at 3
            assert_eq!(bvec.write(3, &data).unwrap(), data.len());
            assert_eq!(bvec.capacity, 1);
            assert_eq!(bvec.size(), 3 + data.len());

            assert_eq!(&bvec[0][..bvec.size()], &[1, 2, 3, 1, 2, 3, 4, 5]);
        }

        {
            let data = [5, 4, 3, 2, 1];
            let mut bvec = BVec::<A>::new();

            // Write at the end of L0-region
            assert_eq!(
                bvec.write((L0_BLOCKS * block::SIZE) as u64, &data).unwrap(),
                data.len()
            );
            // L0_BLOCKS + 1 L1 data block
            assert_eq!(bvec.capacity, L0_BLOCKS + 1);
            assert_eq!(bvec.size(), L0_BLOCKS * block::SIZE + data.len());

            assert_eq!(&bvec[L0_BLOCKS][..data.len()], &data[..]);

            // Write at zero
            assert_eq!(bvec.write(0, &data).unwrap(), data.len());
            assert_eq!(bvec.capacity, L0_BLOCKS + 1);
            assert_eq!(bvec.size(), L0_BLOCKS * block::SIZE + data.len());

            assert_eq!(&bvec[0][..data.len()], &data[..]);

            // Test write crossing L0 block boundary
            assert_eq!(
                bvec.write((block::SIZE - 3) as u64, &data).unwrap(),
                data.len()
            );
            assert_eq!(bvec.capacity, L0_BLOCKS + 1);
            assert_eq!(bvec.size(), L0_BLOCKS * block::SIZE + data.len());

            assert_eq!(&bvec[0][block::SIZE - 3..], &[5, 4, 3]);
            assert_eq!(&bvec[1][..2], &[2, 1]);

            // Test write crossing L0-L1 boundary
            assert_eq!(
                bvec.write((L0_BLOCKS * block::SIZE) as u64 - 2, &data)
                    .unwrap(),
                data.len()
            );
            assert_eq!(bvec.capacity, L0_BLOCKS + 1);
            assert_eq!(bvec.size(), L0_BLOCKS * block::SIZE + data.len());

            assert_eq!(&bvec[L0_BLOCKS - 1][block::SIZE - 2..], &[5, 4]);
            assert_eq!(&bvec[L0_BLOCKS][..data.len()], &[3, 2, 1, 2, 1]);
        }
    }
}
