use std::ops::{Bound, RangeBounds};

use heed::BytesEncode;
use roaring::RoaringBitmap;

use super::{get_first_facet_value, get_highest_level, get_last_facet_value};
use crate::heed_codec::facet::{
    FacetGroupKey, FacetGroupKeyCodec, FacetGroupLazyValueCodec, FacetGroupValueCodec,
};
use crate::heed_codec::BytesRefCodec;
use crate::{CboRoaringBitmapCodec, Result};

/// Find all the document ids for which the given field contains a value contained within
/// the two bounds.
pub fn find_docids_of_facet_within_bounds<'t, BoundCodec>(
    rtxn: &'t heed::RoTxn<'t>,
    db: heed::Database<FacetGroupKeyCodec<BoundCodec>, FacetGroupValueCodec>,
    field_id: u16,
    left: &'t Bound<<BoundCodec as BytesEncode<'t>>::EItem>,
    right: &'t Bound<<BoundCodec as BytesEncode<'t>>::EItem>,
    universe: Option<&RoaringBitmap>,
    docids: &mut RoaringBitmap,
) -> Result<()>
where
    BoundCodec: for<'a> BytesEncode<'a>,
    for<'a> <BoundCodec as BytesEncode<'a>>::EItem: Sized,
{
    let inner;
    let left = match left {
        Bound::Included(left) => {
            inner = BoundCodec::bytes_encode(left).map_err(heed::Error::Encoding)?;
            Bound::Included(inner.as_ref())
        }
        Bound::Excluded(left) => {
            inner = BoundCodec::bytes_encode(left).map_err(heed::Error::Encoding)?;
            Bound::Excluded(inner.as_ref())
        }
        Bound::Unbounded => Bound::Unbounded,
    };
    let inner;
    let right = match right {
        Bound::Included(right) => {
            inner = BoundCodec::bytes_encode(right).map_err(heed::Error::Encoding)?;
            Bound::Included(inner.as_ref())
        }
        Bound::Excluded(right) => {
            inner = BoundCodec::bytes_encode(right).map_err(heed::Error::Encoding)?;
            Bound::Excluded(inner.as_ref())
        }
        Bound::Unbounded => Bound::Unbounded,
    };
    let db = db.remap_types::<FacetGroupKeyCodec<BytesRefCodec>, FacetGroupLazyValueCodec>();
    let mut f = FacetRangeSearch { rtxn, db, field_id, left, right, universe, docids };
    let highest_level = get_highest_level(rtxn, db, field_id)?;

    if let Some(starting_left_bound) =
        get_first_facet_value::<BytesRefCodec, _>(rtxn, db, field_id)?
    {
        let rightmost_bound =
            Bound::Included(get_last_facet_value::<BytesRefCodec, _>(rtxn, db, field_id)?.unwrap()); // will not fail because get_first_facet_value succeeded
        let group_size = usize::MAX;
        f.run(highest_level, starting_left_bound, rightmost_bound, group_size)?;
        Ok(())
    } else {
        Ok(())
    }
}

/// Fetch the document ids that have a facet with a value between the two given bounds
struct FacetRangeSearch<'t, 'b, 'bitmap> {
    rtxn: &'t heed::RoTxn<'t>,
    db: heed::Database<FacetGroupKeyCodec<BytesRefCodec>, FacetGroupLazyValueCodec>,
    field_id: u16,
    left: Bound<&'b [u8]>,
    right: Bound<&'b [u8]>,
    /// The subset of documents ids that are useful for this search.
    /// Great performance optimizations can be achieved by only fetching values matching this subset.
    universe: Option<&'bitmap RoaringBitmap>,
    docids: &'bitmap mut RoaringBitmap,
}

impl<'t> FacetRangeSearch<'t, '_, '_> {
    fn run_level_0(&mut self, starting_left_bound: &'t [u8], group_size: usize) -> Result<()> {
        let left_key =
            FacetGroupKey { field_id: self.field_id, level: 0, left_bound: starting_left_bound };
        let iter = self.db.range(self.rtxn, &(left_key..))?.take(group_size);
        for el in iter {
            let (key, value) = el?;
            // the right side of the iter range is unbounded, so we need to make sure that we are not iterating
            // on the next field id
            if key.field_id != self.field_id {
                return Ok(());
            }
            let should_skip = {
                match self.left {
                    Bound::Included(left) => left > key.left_bound,
                    Bound::Excluded(left) => left >= key.left_bound,
                    Bound::Unbounded => false,
                }
            };
            if should_skip {
                continue;
            }
            let should_stop = {
                match self.right {
                    Bound::Included(right) => right < key.left_bound,
                    Bound::Excluded(right) => right <= key.left_bound,
                    Bound::Unbounded => false,
                }
            };
            if should_stop {
                break;
            }

            if RangeBounds::<&[u8]>::contains(&(self.left, self.right), &key.left_bound) {
                *self.docids |= match self.universe {
                    Some(universe) => CboRoaringBitmapCodec::intersection_with_serialized(
                        value.bitmap_bytes,
                        universe,
                    )?,
                    None => CboRoaringBitmapCodec::deserialize_from(value.bitmap_bytes)?,
                };
            }
        }
        Ok(())
    }

    /// Recursive part of the algorithm for level > 0.
    ///
    /// It works by visiting a slice of a level and checking whether the range asscociated
    /// with each visited element is contained within the bounds.
    ///
    /// 1. So long as the element's range is less than the left bound, we do nothing and keep iterating
    /// 2. If the element's range is fully contained by the bounds, then all of its docids are added to
    ///    the roaring bitmap.
    /// 3. If the element's range merely intersects the bounds, then we call the algorithm recursively
    ///    on the children of the element from the level below.
    /// 4. If the element's range is greater than the right bound, we do nothing and stop iterating.
    ///    Note that the right bound is found through either the `left_bound` of the *next* element,
    ///    or from the `rightmost_bound` argument
    ///
    /// ## Arguments
    /// - `level`: the level being visited
    /// - `starting_left_bound`: the left_bound of the first element to visit
    /// - `rightmost_bound`: the right bound of the last element that should be visited
    /// - `group_size`: the number of elements that should be visited
    fn run(
        &mut self,
        level: u8,
        starting_left_bound: &'t [u8],
        rightmost_bound: Bound<&'t [u8]>,
        group_size: usize,
    ) -> Result<()> {
        if level == 0 {
            return self.run_level_0(starting_left_bound, group_size);
        }

        let left_key =
            FacetGroupKey { field_id: self.field_id, level, left_bound: starting_left_bound };
        let mut iter = self.db.range(self.rtxn, &(left_key..))?.take(group_size);

        // We iterate over the range while keeping in memory the previous value
        let (mut previous_key, mut previous_value) = iter.next().unwrap()?;
        for el in iter {
            let (next_key, next_value) = el?;
            // the right of the iter range is potentially unbounded (e.g. if `group_size` is usize::MAX),
            // so we need to make sure that we are not iterating on the next field id
            if next_key.field_id != self.field_id {
                break;
            }
            // now, do we skip, stop, or visit?
            let should_skip = {
                match self.left {
                    Bound::Included(left) => left >= next_key.left_bound,
                    Bound::Excluded(left) => left >= next_key.left_bound,
                    Bound::Unbounded => false,
                }
            };
            if should_skip {
                previous_key = next_key;
                previous_value = next_value;
                continue;
            }

            // should we stop?
            // We should if the search range doesn't include any
            // element from the previous key or its successors
            let should_stop = {
                match self.right {
                    Bound::Included(right) => right < previous_key.left_bound,
                    Bound::Excluded(right) => right <= previous_key.left_bound,
                    Bound::Unbounded => false,
                }
            };
            if should_stop {
                return Ok(());
            }
            // should we take the whole thing, without recursing down?
            let should_take_whole_group = {
                let left_condition = match self.left {
                    Bound::Included(left) => previous_key.left_bound >= left,
                    Bound::Excluded(left) => previous_key.left_bound > left,
                    Bound::Unbounded => true,
                };
                let right_condition = match self.right {
                    Bound::Included(right) => next_key.left_bound <= right,
                    Bound::Excluded(right) => next_key.left_bound <= right,
                    Bound::Unbounded => true,
                };
                left_condition && right_condition
            };
            if should_take_whole_group {
                *self.docids |= match self.universe {
                    Some(universe) => CboRoaringBitmapCodec::intersection_with_serialized(
                        previous_value.bitmap_bytes,
                        universe,
                    )?,
                    None => CboRoaringBitmapCodec::deserialize_from(previous_value.bitmap_bytes)?,
                };
                previous_key = next_key;
                previous_value = next_value;
                continue;
            }
            // from here, we should visit the children of the previous element and
            // call the function recursively

            let level = level - 1;
            let starting_left_bound = previous_key.left_bound;
            let rightmost_bound = Bound::Excluded(next_key.left_bound);
            let group_size = previous_value.size as usize;

            self.run(level, starting_left_bound, rightmost_bound, group_size)?;

            previous_key = next_key;
            previous_value = next_value;
        }
        // previous_key/previous_value are the last element's key/value

        // now, do we skip, stop, or visit?
        let should_skip = {
            match (self.left, rightmost_bound) {
                (Bound::Included(left), Bound::Included(right)) => left > right,
                (Bound::Included(left), Bound::Excluded(right)) => left >= right,
                (Bound::Excluded(left), Bound::Included(right) | Bound::Excluded(right)) => {
                    left >= right
                }
                (Bound::Unbounded, _) => false,
                (_, Bound::Unbounded) => false, // should never run?
            }
        };
        if should_skip {
            return Ok(());
        }

        // should we stop?
        // We should if the search range doesn't include any
        // element from the previous key or its successors
        let should_stop = {
            match self.right {
                Bound::Included(right) => right < previous_key.left_bound,
                Bound::Excluded(right) => right <= previous_key.left_bound,
                Bound::Unbounded => false,
            }
        };
        if should_stop {
            return Ok(());
        }
        // should we take the whole thing, without recursing down?
        let should_take_whole_group = {
            let left_condition = match self.left {
                Bound::Included(left) => previous_key.left_bound >= left,
                Bound::Excluded(left) => previous_key.left_bound > left,
                Bound::Unbounded => true,
            };
            let right_condition = match (self.right, rightmost_bound) {
                (Bound::Included(right), Bound::Included(rightmost)) => {
                    // we need to stay within the bound ..=right
                    // the element's range goes to ..=righmost
                    // so the element fits entirely within the bound if rightmost <= right
                    rightmost <= right
                }
                (Bound::Included(right), Bound::Excluded(rightmost)) => {
                    // we need to stay within the bound ..=right
                    // the element's range goes to ..righmost
                    // so the element fits entirely within the bound if rightmost <= right
                    rightmost <= right
                }
                (Bound::Excluded(right), Bound::Included(rightmost)) => {
                    // we need to stay within the bound ..right
                    // the element's range goes to ..=righmost
                    // so the element fits entirely within the bound if rightmost < right
                    rightmost < right
                }
                (Bound::Excluded(right), Bound::Excluded(rightmost)) => {
                    // we need to stay within the bound ..right
                    // the element's range goes to ..righmost
                    // so the element fits entirely within the bound if rightmost <= right
                    rightmost <= right
                }
                (Bound::Unbounded, _) => {
                    // we need to stay within the bound ..inf
                    // so the element always fits entirely within the bound
                    true
                }
                (_, Bound::Unbounded) => {
                    // we need to stay within a finite bound
                    // but the element's range goes to ..inf
                    // so the element never fits entirely within the bound
                    false
                }
            };
            left_condition && right_condition
        };
        if should_take_whole_group {
            *self.docids |= match self.universe {
                Some(universe) => CboRoaringBitmapCodec::intersection_with_serialized(
                    previous_value.bitmap_bytes,
                    universe,
                )?,
                None => CboRoaringBitmapCodec::deserialize_from(previous_value.bitmap_bytes)?,
            };
        } else {
            let level = level - 1;
            let starting_left_bound = previous_key.left_bound;
            let group_size = previous_value.size as usize;

            self.run(level, starting_left_bound, rightmost_bound, group_size)?;
        }

        Ok(())
    }
}
