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())
    }
}

/// Byte Vector implemented as tree of pages.
///
/// root and len change as the TreeVec increases in size.
pub struct TreeVec<'a> {
    pub root: u64,
    pub len: u64,
    pub ps: &'a mut PageSet,
    writing: bool,
    pub len_changed: bool,
    pub root_changed: bool,
}

impl<'a> TreeVec<'a> {
    /// Create a TreeVec to access contents( read/write ).
    pub fn new(root: u64, len: u64, ps: &'a mut PageSet) -> Self {
        Self {
            root,
            len,
            ps,
            writing: false,
            len_changed: false,
            root_changed: false,
        }
    }

    /// Increase length to specified value ( which must be >= current len ).
    pub fn _resize(&mut self, len: u64) {
        assert!(len >= self.len);
        self.adjust_len(len);
    }

    /// Reduce length to specified value.
    pub fn _truncate(&mut self, _len: u64) {
        todo!()
    }

    /// Write bytes at specified index.
    pub fn write(&mut self, mut ix: u64, user_data: &[u8]) {
        self.writing = true;
        let mut todo = user_data.len();

        let levels = self.adjust_len(ix + todo as u64);

        let mut done = 0;
        while todo > 0 {
            let rpp = PAGE_SIZE;
            let mut pnum = self.root;
            let mut off = ix;
            if levels > 0 {
                pnum = self.get_page(levels, pnum, ix / rpp);
                off = ix % rpp;
            }

            let (mut data, _changed) = self.ps.load(pnum);
            let md = Arc::make_mut(&mut data);

            let off = off as usize;
            let mut amount = PAGE_SIZE as usize - off;
            if amount > todo {
                amount = todo;
            }

            let end = off + amount;
            if end > md.len() {
                md.resize(end, 0);
            }
            let loc = &mut md[off..end];
            loc.copy_from_slice(&user_data[done..done + amount]);
            ix += amount as u64;
            done += amount;
            todo -= amount;
            self.ps.note(pnum, data, true);
        }
        self.writing = false;
    }

    /// Read bytes from specified index.
    ///
    /// If un-written bytes are detected, reading stops. Result is number of bytes read.
    /// However un-written bytes can also be read as zeroes. Often user_data is initialised to zero, so result is same.
    pub fn read(&mut self, mut ix: u64, user_data: &mut [u8]) -> usize {
        let mut todo = user_data.len();
        if ix + todo as u64 > self.len {
            return 0;
        }
        let levels = self.levels(self.len);

        let mut done = 0;
        while todo > 0 {
            let rpp = PAGE_SIZE;
            let mut pnum = self.root;
            let mut off = ix;
            if levels > 0 {
                pnum = self.get_page(levels, pnum, ix / rpp);
                if pnum == 0 {
                    return done;
                }
                off = ix % rpp;
            }

            let (data, changed) = self.ps.load(pnum);

            let off = off as usize;
            let mut amount = PAGE_SIZE as usize - off;
            if amount > todo {
                amount = todo;
            }

            let end = off + amount;
            if end > data.len() {
                return done;
            }
            let loc = &data[off..end];

            user_data[done..done + amount].copy_from_slice(loc);
            ix += amount as u64;
            done += amount;
            todo -= amount;
            self.ps.note(pnum, data, changed);
        }
        done
    }

    /// Get a page at specified level.
    fn get_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.get_page(level - 1, page, x);
        }
        if page == 0 {
            return 0;
        }
        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_len(&mut self, size: u64) -> u8 {
        let mut level = self.levels(self.len);
        if size > self.len {
            let new_level = self.levels(size);

            while level < new_level {
                self.inc_level();
                level += 1;
            }
            self.len = size;
            self.len_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 TreeVec of specified size.
    fn levels(&self, size: u64) -> u8 {
        let rpp = PAGE_SIZE;
        let base = PAGE_SIZE / 8; // Number of child pages per page.
        let mut pages = size.div_ceil(rpp);
        if pages <= 1 {
            return 0;
        }
        let mut result = 1;
        while pages > base {
            pages = pages.div_ceil(base);
            result += 1;
        }
        result
    }

    /// Page interface : get a page number. self.ps can then be used to load or note the page data,.
    pub fn _get_page_num(&mut self, pix: u64, create: bool) -> u64 {
        self.writing = create;
        let levels = self.levels(self.len);
        let result = if levels > 0 {
            self.get_page(levels - 1, self.root, pix)
        } else {
            assert!(pix == 0);
            self.root
        };
        self.writing = false;
        result
    }

    pub fn read_u64(&mut self, off: u64) -> u64 {
        let mut buf = [0u8; 8];
        self.read(off, &mut buf);
        u64::from_le_bytes(buf)
    }

    pub fn write_u64(&mut self, off: u64, val: u64) {
        let mut buf = [0u8; 8];
        buf.copy_from_slice(&val.to_le_bytes());
        self.write(off, &buf);
    }
}

#[cfg(test)]
pub fn test_tv(root: u64, len: u64, ps: &mut PageSet) {
    let mut tv = TreeVec::new(root, len, ps);

    let data = b"hello";
    tv.write(5 * 25, data);

    let mut buf = [0; 5];
    tv.read(5 * 25, &mut buf);
    println!("rd={:?}", tos(&buf));
    assert_eq!(data, &buf);

    let data = b"there";
    tv.write(5 * 53, data);

    let mut buf = [0; 5];
    tv.read(5 * 53, &mut buf);
    println!("rd={:?}", tos(&buf));
    assert_eq!(data, &buf);

    let data = b"georg";
    tv.write(5 * 81, data);

    let mut buf = [0; 5];
    tv.read(5 * 81, &mut buf);
    println!("rd={:?}", tos(&buf));
    assert_eq!(data, &buf);

    println!("tv root={} len={}", tv.root, tv.len);
}
