// Copyright 2015 Pierre-Étienne Meunier and Florent Becker. See the
// COPYRIGHT file at the top-level directory of this distribution and
// at http://pijul.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use crate::skiplist;
use crate::skiplist::{record_size, FIRST_BINDING_SIZE};
use crate::transaction::{Cow, Env, Exclusive};
use crate::{rc, Cursor, Db, Error, MutTxn, PageItem, PageT, Representable, PAGE_SIZE_U16};
use rand::Rng;
use std;
use std::borrow::Borrow;

impl<E: Borrow<Env<Exclusive>>, T> MutTxn<E, T> {
    /// Insert a binding to a database, returning `false` if and only
    /// if the exact same binding (key *and* value) was already in the database.
    pub fn put<R: Rng, K: Representable, V: Representable>(
        &mut self,
        rng: &mut R,
        db: &mut Db<K, V>,
        key: K,
        value: V,
    ) -> Result<bool, Error> {
        let mut cursor = Cursor::new();
        if cursor.set(self, db, Some((key, Some(value))))? {
            return Ok(false);
        }
        // Now insert in the leaf and split all pages that need to be split.

        let leaf_pointer = cursor.pointer;

        // Split all pages that need it.
        //
        // This loop uses unsafe slices, but the memory management is
        // simple: all pointers used in the loop are valid until we
        // explicitely call free, after the loop.
        let last_allocated = self.put_cascade(rng, &mut cursor, key, value)?;

        // Now, cursor_pointer is the reference to a page that doesn't
        // need to split anymore (possibly the new root).

        // Free all split pages.
        self.free_split_pages(rng, &mut cursor, leaf_pointer)?;

        // `put_cascade` above only returns after the insertion of a
        // (key, value) with right child `inserted_right_child`. For
        // `put_update_split_root` to update the reference properly in
        // the pages above the insertion, we need to update that child
        // here (it might be 0 if we're still at the leaf).

        // Has the root split?
        self.update_split_root(rng, &mut cursor, last_allocated, db)?;
        debug!("put done");
        Ok(true)
    }

    fn put_cascade<R: Rng, K: Representable, V: Representable>(
        &mut self,
        rng: &mut R,
        cursor: &mut Cursor,
        key: K,
        value: V,
    ) -> Result<u64, Error> {
        let mut last_allocated = 0;
        let mut inserted_right_child = 0;
        let mut inserted_key = key;
        let mut inserted_value = value;
        let mut inserted_left_child = 0u64;

        loop {
            let mut page = self.load_cow_page(cursor.current().page)?;

            debug!("page: {:?}", page);

            // If we're "CoWing" the first double-pointed page on the stack, decrease its RC.
            if cursor.pointer == cursor.first_rc_level {
                let page = cursor.current().page;
                let rc = rc(self, page)?;
                if rc >= 2 {
                    self.set_rc(rng, page, rc - 1)?;
                }
            }

            if skiplist::can_alloc(&page, inserted_key, inserted_value) != 0 {
                self.put_can_alloc(
                    rng,
                    page,
                    cursor,
                    &mut inserted_key,
                    &mut inserted_left_child,
                    &mut inserted_right_child,
                    &mut inserted_value,
                    &mut last_allocated,
                )?;
                break;
            } else {
                // We cannot allocate, we need to split this page

                debug!("cannot alloc");
                self.put_split(
                    rng,
                    &mut page,
                    cursor,
                    &mut inserted_key,
                    &mut inserted_left_child,
                    &mut inserted_right_child,
                    &mut inserted_value,
                )?;

                cursor.pointer -= 1;
                if cursor.pointer == 0 {
                    // We've just split the root! we need to
                    // allocate a page to collect the two sides of
                    // the split.
                    last_allocated = self.split_root(
                        rng,
                        cursor,
                        inserted_left_child,
                        inserted_right_child,
                        inserted_key,
                        inserted_value,
                    )?;
                    break;
                }
            }
        }
        Ok(last_allocated)
    }

    fn put_can_alloc<R: Rng, K: Representable, V: Representable>(
        &mut self,
        rng: &mut R,
        mut page: Cow,
        cursor: &mut Cursor,
        inserted_key: &mut K,
        inserted_left_child: &mut u64,
        inserted_right_child: &mut u64,
        inserted_value: &mut V,
        last_allocated: &mut u64,
    ) -> Result<(), Error> {
        debug!("can alloc");
        // We can allocate. Do we need to duplicate this
        // page (because of RC or first_free)?
        let needs_dup = skiplist::first_free(&page) + record_size(*inserted_key, *inserted_value)
            > PAGE_SIZE_U16;

        match page {
            Cow::MutPage(ref mut page) if cursor.pointer < cursor.first_rc_level && !needs_dup => {
                debug!("does not need dup");

                // No, there will be no other references to
                // this page after this method (`put`)
                // returns. Here, simply insert.

                // Update the right child from the previous split
                // (if any, `inserted_left_child` is 0 else).
                debug!("page right child: {:?}", cursor.current()[0]);
                self.set_right_child(page, cursor.current()[0], *inserted_left_child);
                self.skiplist_insert_after_cursor(
                    page,
                    rng,
                    &mut cursor.current_mut().cursor,
                    *inserted_key,
                    *inserted_value,
                    *inserted_right_child,
                );
                *last_allocated = page.page_offset();
            }
            _ => {
                debug!(
                    "needs dup {:?} {:?} {:?}",
                    cursor.pointer, cursor.first_rc_level, needs_dup
                );
                // Yes, we do need to duplicate this
                // page. Propagate RC to pages
                // below.
                let page = self.load_cow_page(cursor.current().page)?;

                // Iterate over the page, incorporating the new element.
                let it = Iter {
                    iter: skiplist::iter_all(&page).peekable(),
                    off: page.offset(cursor.current()[0] as isize),
                    next_is_extra: false,
                    extra_left: *inserted_left_child,
                    extra: (
                        Some((*inserted_key, *inserted_value)),
                        *inserted_right_child,
                    ),
                };

                // Increase the RC of everyone, but the left
                // side of the split (which was just
                // allocated now).
                let incr_rc = cursor.pointer >= cursor.first_rc_level;
                *last_allocated = self.spread_on_one_page(
                    rng,
                    it.map(|(b, c)| {
                        let fresh = c == *inserted_left_child || c == *inserted_right_child;
                        PageItem {
                            ptr: std::ptr::null(),
                            kv: b,
                            right: c,
                            incr_right_rc: incr_rc && !fresh,
                            incr_kv_rc: incr_rc,
                        }
                    }),
                )?;

                if cursor.pointer < cursor.first_rc_level {
                    unsafe {
                        // free of the current page.
                        self.free_page(cursor.current().page)
                    }
                }
            }
        }
        Ok(())
    }

    fn put_split<R: Rng, K: Representable, V: Representable>(
        &mut self,
        rng: &mut R,
        page: &mut Cow,
        cursor: &mut Cursor,
        inserted_key: &mut K,
        inserted_left_child: &mut u64,
        inserted_right_child: &mut u64,
        inserted_value: &mut V,
    ) -> Result<(), Error> {
        let it = Iter {
            iter: skiplist::iter_all(page).peekable(),
            off: page.offset(cursor.current()[0] as isize),
            next_is_extra: false,
            extra_left: *inserted_left_child,
            extra: (
                Some((*inserted_key, *inserted_value)),
                *inserted_right_child,
            ),
        };
        // We're going to add a page, hence
        // FIRST_BINDING_SIZE more bytes will be
        // needed.
        let total = skiplist::occupied(page)
            + record_size(*inserted_key, *inserted_value)
            + FIRST_BINDING_SIZE;
        let (left, right, sep_key, sep_value) = if cursor.pointer >= cursor.first_rc_level {
            let it = it.map(|(b, c)| {
                if c == *inserted_left_child || c == *inserted_right_child {
                    PageItem {
                        ptr: std::ptr::null(),
                        kv: b,
                        right: c,
                        incr_right_rc: false,
                        incr_kv_rc: true,
                    }
                } else {
                    PageItem {
                        ptr: std::ptr::null(),
                        kv: b,
                        right: c,
                        incr_right_rc: true,
                        incr_kv_rc: true,
                    }
                }
            });
            self.spread_on_two_pages(rng, it, total)?
        } else {
            let it = it.map(|(b, c)| PageItem {
                ptr: std::ptr::null(),
                kv: b,
                right: c,
                incr_right_rc: false,
                incr_kv_rc: false,
            });
            self.spread_on_two_pages(rng, it, total)?
        };
        // Extend the lifetime of "inserted_key".
        *inserted_key = sep_key;
        *inserted_left_child = left;
        *inserted_right_child = right;
        *inserted_value = sep_value;
        Ok(())
    }

    fn update_put<R: Rng, K: Representable, V: Representable>(
        &mut self,
        rng: &mut R,
        cursor: &mut Cursor,
    ) -> Result<(), Error> {
        let mut page = self.load_cow_page(cursor.current().page)?;
        match page {
            Cow::MutPage(ref mut p) if cursor.pointer < cursor.first_rc_level => {
                self.set_right_child(p, cursor.current()[0], cursor[cursor.pointer + 1].page);
                cursor.current_mut().page = p.offset;
            }
            Cow::MutPage(_) | Cow::Page(_) => {
                // Here, the page needs to be CoWed.
                let incr_rc = cursor.pointer >= cursor.first_rc_level;
                let current_ptr = page.offset(cursor.current()[0] as isize);
                let new_page = self.spread_on_one_page::<_, K, V, _>(
                    rng,
                    skiplist::iter_all(&page).map(|(off, b, c)| {
                        // Both pages will survive, increase reference count.
                        let mut item = PageItem {
                            ptr: std::ptr::null(),
                            kv: b,
                            right: c,
                            incr_kv_rc: incr_rc,
                            incr_right_rc: incr_rc,
                        };
                        if off == current_ptr {
                            item.right = cursor[cursor.pointer + 1].page;
                            item.incr_right_rc = false // fresh page.
                        }
                        item
                    }),
                )?;
                // If nothing points to the page anymore, free.
                if cursor.pointer < cursor.first_rc_level {
                    unsafe {
                        debug!("line {:?}, freeing {:?}", line!(), page.page_offset());
                        self.free_page(page.page_offset())
                    }
                }
                debug!("NEW_PAGE = {:?}", new_page);
                cursor.current_mut().page = new_page;
            }
        }
        Ok(())
    }

    fn free_split_pages<R: Rng>(
        &mut self,
        rng: &mut R,
        cursor: &mut Cursor,
        leaf_pointer: usize,
    ) -> Result<(), Error> {
        debug!(
            "freeing pages from {:?} to {:?} {:?} (inclusive)",
            cursor.pointer + 1,
            leaf_pointer,
            cursor.first_rc_level
        );
        let last_free = std::cmp::min(leaf_pointer + 1, cursor.first_rc_level);
        let first_free = std::cmp::min(cursor.pointer + 1, last_free);
        for page_cursor in &cursor.stack[first_free..last_free] {
            // free the page that was split.
            let rc = rc(self, page_cursor.page)?;
            if rc <= 1 {
                debug!("Freeing {:?}", page_cursor.page);
                self.remove_rc(rng, page_cursor.page)?;
                unsafe { self.free_page(page_cursor.page) }
            } else {
                debug!(
                    "Decrease RC for page {:?}, new rc: {:?}",
                    page_cursor.page,
                    rc - 1
                );
                self.set_rc(rng, page_cursor.page, rc - 1)?;
            }
        }
        debug!("freed");
        Ok(())
    }

    fn update_split_root<R: Rng, K: Representable, V: Representable>(
        &mut self,
        rng: &mut R,
        cursor: &mut Cursor,
        last_allocated: u64,
        db: &mut Db<K, V>,
    ) -> Result<(), Error> {
        if cursor.pointer > 0 {
            cursor.stack[cursor.pointer].page = last_allocated;
            cursor.pointer -= 1;
            // The root has not split, the splits stopped earlier. There are remaining pages.
            while cursor.pointer > 0 {
                if cursor.pointer == cursor.first_rc_level {
                    let page = cursor.current().page;
                    let rc = rc(self, page)?;
                    if rc >= 2 {
                        self.set_rc(rng, page, rc - 1)?;
                    }
                }
                self.update_put::<R, K, V>(rng, cursor)?;
                cursor.pointer -= 1;
            }

            db.0 = cursor[1].page
        } else {
            // The root has split, no need to do anything.
            debug!("the root has split");
            db.0 = last_allocated
        }
        Ok(())
    }
}

/// Iterator on a page, plus a middle element immediately after `off`,
/// with a fixed left and right child.
struct Iter<'a, P: PageT + Sized + 'a, K: Representable, V: Representable> {
    iter: std::iter::Peekable<skiplist::Iter<'a, P, K, V>>,
    off: *const u8,
    next_is_extra: bool,
    extra_left: u64,
    extra: (Option<(K, V)>, u64),
}

impl<'a, P: PageT + Sized + 'a, K: Representable, V: Representable> Iterator for Iter<'a, P, K, V> {
    type Item = (Option<(K, V)>, u64);

    fn next(&mut self) -> Option<Self::Item> {
        if self.next_is_extra {
            debug!("next_is_extra");
            self.next_is_extra = false;
            self.off = std::ptr::null();
            Some(self.extra)
        } else {
            debug!("not next_is_extra");

            if let Some((a, b, d)) = self.iter.next() {
                // if a == self.off, the extra element is to be
                // inserted immediately after this element.  Set
                // next_is_extra, and replace the current element's
                // right child with extra's left child.
                debug!("{:?}", a);
                if a == self.off {
                    self.next_is_extra = true;
                    Some((b, self.extra_left))
                } else {
                    Some((b, d))
                }
            } else {
                debug!("finished, extra = {:?}", self.off);
                None
            }
        }
    }
}