use crate::*;

/// Parent page is list of page numbers.
struct ParentPage {
    data: Data,
}

impl ParentPage {
    fn new() -> Self {
        Self {
            data: Arc::new(Vec::new()),
        }
    }

    fn from(data: Data) -> Self {
        Self { data }
    }

    fn take(self) -> Data {
        self.data
    }

    /// Assign page number.
    fn assign(&mut self, ix: usize, pnum: u64) {
        let data = Arc::make_mut(&mut self.data);
        let off = ix * 8;
        let end = off + 8;
        if end > data.len() {
            data.resize(end, 0);
        }
        let loc = &mut data[off..end];
        loc.copy_from_slice(&pnum.to_le_bytes());
    }

    /// Fetch page number.
    fn fetch(&self, ix: usize) -> u64 {
        let off = ix * 8;
        let end = off + 8;
        if end > self.data.len() {
            return 0;
        }
        let loc = &self.data[off..end];
        u64::from_le_bytes(loc.try_into().unwrap())
    }
}

/// Tree of pages.
pub struct PageTree<'a> {
    pub root: u64,
    pub count: u64,
    pub ps: &'a mut PageSet,
    writing: bool,
    pub count_changed: bool,
    pub root_changed: bool,
}

impl<'a> PageTree<'a> {
    /// Create or restore a PageTree.
    pub fn new(root: u64, count: u64, ps: &'a mut PageSet) -> Self {
        Self {
            root,
            count,
            ps,
            writing: false,
            count_changed: false,
            root_changed: false,
        }
    }

    /// Increase page count to specified value ( which must be >= current count ).
    pub fn resize(&mut self, count: u64) {
        assert!(count >= self.count);
        self.adjust_count(count);
    }

    /// Get a page number.
    pub fn get(&mut self, pix: u64, create: bool) -> u64 {
        assert!(pix < self.count);
        self.writing = create;
        let levels = self.levels(self.count);
        let result = if levels > 0 {
            self.page(levels - 1, self.root, pix)
        } else {
            assert!(pix == 0);
            self.root
        };
        self.writing = false;
        result
    }

    /// Free all the pages. PageTree is no longer useable.
    pub fn drop_pages(&mut self) {
        let mut count = self.count;
        let mut level = self.levels(self.count);
        while level > 0 {
            let base = PAGE_SIZE / 8; // Number of child pages per page.
            for pix in 0..count {
                let pnum = self.page(level, self.root, pix);
                if pnum != 0 {
                    self.ps.free_page(pnum);
                }
            }
            level -= 1;
            count = count.div_ceil(base);
        }
        self.ps.free_page(self.root);
        self.root = 0;
        self.count = 0;
    }

    /// Get a page number at specified level.
    fn page(&mut self, level: u8, mut page: u64, mut ix: u64) -> u64 {
        let base = PAGE_SIZE / 8;
        if level > 1 {
            let x = ix / base;
            ix %= base;
            page = self.page(level - 1, page, x);
        }
        self.get_child_page(page, ix as usize)
    }

    /// Get a child page at specified index.
    fn get_child_page(&mut self, page: u64, ix: usize) -> u64 {
        let (data, mut changed) = self.ps.load(page);
        let mut pp = ParentPage::from(data);
        let mut result = pp.fetch(ix);
        if result == 0 {
            if !self.writing {
                return 0;
            }
            result = self.ps.new_page();
            pp.assign(ix, result);
            changed = true;
        }
        self.ps.note(page, pp.take(), changed);
        result
    }

    /// Increase the TreeVec len to specified size, returns the number of levels.
    fn adjust_count(&mut self, count: u64) -> u8 {
        let mut level = self.levels(self.count);
        if count > self.count {
            let new_level = self.levels(count);
            while level < new_level {
                self.inc_level();
                level += 1;
            }
            self.count = count;
            self.count_changed = true;
        }
        level
    }

    /// Increase the TreeVec level by creating a new root. The old root is stored in the new root at position 0.
    fn inc_level(&mut self) {
        let new_root = self.ps.new_page();
        let mut pp = ParentPage::new();
        pp.assign(0, self.root);
        self.ps.note(new_root, pp.take(), true);
        self.root = new_root;
        self.root_changed = true;
    }

    /// Calculate number of extra levels needed for PageTree of specified count.
    fn levels(&self, mut count: u64) -> u8 {
        let base = PAGE_SIZE / 8; // Number of child pages per page.
        if count <= 1 {
            return 0;
        }
        let mut result = 1;
        while count > base {
            count = count.div_ceil(base);
            result += 1;
        }
        result
    }
}
