use std::cmp::Ordering;

use ahash::{AHashMap, AHashSet};

use crate::{ReadTransaction, Store, ValueKey};

use super::{Comparator, ResultSet, SortedResultSet};

pub struct Pagination {
    requested_position: i32,
    position: i32,
    limit: usize,
    anchor: u32,
    anchor_offset: i32,
    has_anchor: bool,
    anchor_found: bool,
    ids: Vec<u64>,
    prefix_key: Option<ValueKey>,
    prefix_unique: bool,
}

impl ReadTransaction<'_> {
    #[maybe_async::maybe_async]
    pub async fn sort(
        &mut self,
        result_set: ResultSet,
        mut comparators: Vec<Comparator>,
        mut paginate: Pagination,
    ) -> crate::Result<SortedResultSet> {
        if comparators.len() == 1 && !paginate.prefix_unique {
            match comparators.pop().unwrap() {
                Comparator::Field { field, ascending } => {
                    let mut results = result_set.results;

                    self.sort_index(
                        result_set.account_id,
                        result_set.collection,
                        field,
                        ascending,
                        |_, document_id| {
                            !results.remove(document_id) || paginate.add(0, document_id)
                        },
                    )
                    .await?;

                    // Add remaining items not present in the index
                    if !results.is_empty() && !paginate.is_full() {
                        for document_id in results {
                            if !paginate.add(0, document_id) {
                                break;
                            }
                        }
                    }
                }
                Comparator::DocumentSet { set, ascending } => {
                    let in_set = &result_set.results & &set;
                    let not_in_set = &result_set.results ^ &in_set;
                    let sets = if ascending {
                        [in_set, not_in_set]
                    } else {
                        [not_in_set, in_set]
                    };
                    'outer: for set in sets {
                        for document_id in set {
                            if !paginate.add(0, document_id) {
                                break 'outer;
                            }
                        }
                    }
                }
            }

            // Obtain prefixes
            let prefix_key = paginate.prefix_key.take();
            let mut sorted_results = paginate.build();
            if let Some(prefix_key) = prefix_key {
                for id in sorted_results.ids.iter_mut() {
                    if let Some(prefix_id) = self
                        .get_value::<u32>(prefix_key.with_document_id(*id as u32))
                        .await?
                    {
                        *id |= (prefix_id as u64) << 32;
                    }
                }
            }

            Ok(sorted_results)
        } else {
            let mut sorted_ids = AHashMap::with_capacity(paginate.limit);

            for (pos, comparator) in comparators.into_iter().take(4).enumerate() {
                match comparator {
                    Comparator::Field { field, ascending } => {
                        let mut results = result_set.results.clone();
                        let mut prev_data = vec![];
                        let mut has_grouped_ids = false;
                        let mut idx = 0;

                        self.refresh_if_old().await?;
                        self.sort_index(
                            result_set.account_id,
                            result_set.collection,
                            field,
                            ascending,
                            |data, document_id| {
                                if results.remove(document_id) {
                                    debug_assert!(!data.is_empty());

                                    if data != prev_data {
                                        idx += 1;
                                        prev_data = data.to_vec();
                                    } else {
                                        has_grouped_ids = true;
                                    }

                                    sorted_ids.entry(document_id).or_insert([0u32; 4])[pos] = idx;

                                    !results.is_empty()
                                } else {
                                    true
                                }
                            },
                        )
                        .await?;

                        // Add remaining items not present in the index
                        if !results.is_empty() {
                            idx += 1;
                            for document_id in results {
                                sorted_ids.entry(document_id).or_insert([0u32; 4])[pos] = idx;
                            }
                        }

                        if !has_grouped_ids {
                            // If we are sorting by multiple fields and we don't have grouped ids, we can
                            // stop here
                            break;
                        }
                    }
                    Comparator::DocumentSet { set, ascending } => {
                        let in_set = &result_set.results & &set;
                        let not_in_set = &result_set.results ^ &in_set;
                        let sets = if ascending {
                            [(in_set, 0), (not_in_set, 1)]
                        } else {
                            [(not_in_set, 0), (in_set, 1)]
                        };

                        for (document_ids, idx) in sets {
                            for document_id in document_ids {
                                sorted_ids.entry(document_id).or_insert([0u32; 4])[pos] = idx;
                            }
                        }
                    }
                }
            }

            let mut seen_prefixes = AHashSet::new();
            let mut sorted_ids = sorted_ids.into_iter().collect::<Vec<_>>();
            sorted_ids.sort_by(|a, b| match a.1.cmp(&b.1) {
                Ordering::Equal => a.0.cmp(&b.0),
                other => other,
            });
            for (document_id, _) in sorted_ids {
                // Obtain document prefixId
                let prefix_id = if let Some(prefix_key) = &paginate.prefix_key {
                    if let Some(prefix_id) = self
                        .get_value(prefix_key.with_document_id(document_id))
                        .await?
                    {
                        if paginate.prefix_unique && !seen_prefixes.insert(prefix_id) {
                            continue;
                        }
                        prefix_id
                    } else {
                        // Document no longer exists?
                        continue;
                    }
                } else {
                    0
                };

                // Add document to results
                if !paginate.add(prefix_id, document_id) {
                    break;
                }
            }

            Ok(paginate.build())
        }
    }
}

impl Store {
    pub async fn sort(
        &self,
        result_set: ResultSet,
        comparators: Vec<Comparator>,
        mut paginate: Pagination,
    ) -> crate::Result<SortedResultSet> {
        paginate.limit = match (result_set.results.len(), paginate.limit) {
            (0, _) => {
                return Ok(SortedResultSet {
                    position: paginate.position,
                    ids: vec![],
                    found_anchor: true,
                });
            }
            (_, 0) => result_set.results.len() as usize,
            (a, b) => std::cmp::min(a as usize, b),
        };

        #[cfg(feature = "is_async")]
        {
            self.read_transaction()
                .await?
                .sort(result_set, comparators, paginate)
                .await
        }

        #[cfg(feature = "is_sync")]
        {
            let mut trx = self.read_transaction()?;
            self.spawn_worker(move || trx.sort(result_set, comparators, paginate))
                .await
        }
    }
}

impl Pagination {
    pub fn new(
        limit: usize,
        position: i32,
        anchor: Option<u32>,
        anchor_offset: i32,
        prefix_key: Option<ValueKey>,
        prefix_unique: bool,
    ) -> Self {
        let (has_anchor, anchor) = anchor.map(|anchor| (true, anchor)).unwrap_or((false, 0));

        Self {
            requested_position: position,
            position,
            limit,
            anchor,
            anchor_offset,
            has_anchor,
            anchor_found: false,
            ids: Vec::with_capacity(limit),
            prefix_key,
            prefix_unique,
        }
    }

    pub fn add(&mut self, prefix_id: u32, document_id: u32) -> bool {
        let id = ((prefix_id as u64) << 32) | document_id as u64;

        // Pagination
        if !self.has_anchor {
            if self.position >= 0 {
                if self.position > 0 {
                    self.position -= 1;
                } else {
                    self.ids.push(id);
                    if self.ids.len() == self.limit {
                        return false;
                    }
                }
            } else {
                self.ids.push(id);
            }
        } else if self.anchor_offset >= 0 {
            if !self.anchor_found {
                if document_id != self.anchor {
                    return true;
                }
                self.anchor_found = true;
            }

            if self.anchor_offset > 0 {
                self.anchor_offset -= 1;
            } else {
                self.ids.push(id);
                if self.ids.len() == self.limit {
                    return false;
                }
            }
        } else {
            self.anchor_found = document_id == self.anchor;
            self.ids.push(id);

            if self.anchor_found {
                self.position = self.anchor_offset;
                return false;
            }
        }

        true
    }

    pub fn is_full(&self) -> bool {
        self.ids.len() == self.limit
    }

    pub fn build(self) -> SortedResultSet {
        let mut result = SortedResultSet {
            ids: self.ids,
            position: 0,
            found_anchor: !self.has_anchor || self.anchor_found,
        };

        if result.found_anchor {
            if !self.has_anchor && self.requested_position >= 0 {
                result.position = if self.position == 0 {
                    self.requested_position
                } else {
                    0
                };
            } else if self.position >= 0 {
                result.position = self.position;
            } else {
                let position = self.position.unsigned_abs() as usize;
                let start_offset = if position < result.ids.len() {
                    result.ids.len() - position
                } else {
                    0
                };
                result.position = start_offset as i32;
                let end_offset = if self.limit > 0 {
                    std::cmp::min(start_offset + self.limit, result.ids.len())
                } else {
                    result.ids.len()
                };

                result.ids = result.ids[start_offset..end_offset].to_vec()
            }
        }

        result
    }
}