use anyhow::Result;

use super::{ShardReader, ShardReaderStat};
use crate::{reader::query_info::QueryInfo, utils::key::ShardingKeyRange};

pub const DEFAULT_SHARD_SAMPLE_COUNT: usize = 6;
pub const DEFAULT_FRAG_SAMPLE_COUNT: usize = 8;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct EstimateSampleCounts {
    pub shard_sample_count: usize,
    pub frag_sample_count: usize,
}

impl Default for EstimateSampleCounts {
    fn default() -> Self {
        Self {
            shard_sample_count: DEFAULT_SHARD_SAMPLE_COUNT,
            frag_sample_count: DEFAULT_FRAG_SAMPLE_COUNT,
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct ShardEstimatePopulation {
    pub total_frags: usize,
    pub population_live_docs: u64,
}

#[derive(Clone, Debug, Default)]
pub struct ShardEstimateSample {
    pub available_frags: usize,
    pub sampled_frags: usize,
    pub population_live_docs: u64,
    pub sampled_live_docs: u64,
    pub sampled_match_count: u64,
}

#[derive(Clone, Debug)]
struct FragPopulationSummary {
    index: usize,
    estimated_live_docs: u64,
}

#[derive(Clone, Debug, Default)]
struct PopulationSummary {
    total_frags: usize,
    candidates: Vec<FragPopulationSummary>,
    population_live_docs: u64,
}

impl ShardReader {
    pub fn estimate_population(
        &self,
        read_ranges: &[ShardingKeyRange],
    ) -> Result<ShardEstimatePopulation> {
        let population = self.build_population_summary(read_ranges)?;
        Ok(ShardEstimatePopulation {
            total_frags: population.total_frags,
            population_live_docs: population.population_live_docs,
        })
    }

    pub fn sample_count_with_query_info(
        &self,
        query_info: &QueryInfo,
    ) -> Result<ShardEstimateSample> {
        let population = self.build_population_summary(query_info.read_ranges())?;
        if population.population_live_docs == 0 || population.candidates.is_empty() {
            return Ok(ShardEstimateSample {
                available_frags: population.total_frags,
                ..Default::default()
            });
        }

        let mut candidates = population.candidates;
        select_top_live_doc_candidates(
            &mut candidates,
            self.estimate_sample_counts.frag_sample_count,
        );

        let mut search_stat = ShardReaderStat::default();
        let mut sampled_live_docs = 0_u64;
        let mut sampled_match_count = 0_u64;
        let mut sampled_frags = 0_usize;
        for candidate in candidates {
            let frag_reader = &self.frag_readers[candidate.index];
            let match_count =
                frag_reader.search_with_count_collector(query_info, &mut search_stat)?;
            sampled_frags += 1;
            sampled_live_docs += candidate.estimated_live_docs;
            sampled_match_count += match_count;
        }

        Ok(ShardEstimateSample {
            available_frags: population.total_frags,
            sampled_frags,
            population_live_docs: population.population_live_docs,
            sampled_live_docs,
            sampled_match_count,
        })
    }

    fn build_population_summary(
        &self,
        read_ranges: &[ShardingKeyRange],
    ) -> Result<PopulationSummary> {
        let mut population = PopulationSummary {
            total_frags: self.frag_readers.len(),
            ..Default::default()
        };
        for (index, frag_reader) in self.frag_readers.iter().enumerate() {
            let estimated_live_docs =
                frag_reader.estimated_live_doc_count_in_ranges(read_ranges)?;
            if estimated_live_docs == 0 {
                continue;
            }
            population.population_live_docs += estimated_live_docs;
            population.candidates.push(FragPopulationSummary {
                index,
                estimated_live_docs,
            });
        }
        Ok(population)
    }
}

fn select_top_live_doc_candidates(
    candidates: &mut Vec<FragPopulationSummary>,
    frag_sample_count: usize,
) {
    candidates.sort_unstable_by(|left, right| {
        right
            .estimated_live_docs
            .cmp(&left.estimated_live_docs)
            .then_with(|| left.index.cmp(&right.index))
    });
    candidates.truncate(frag_sample_count);
}

#[cfg(test)]
mod tests {
    use super::{DEFAULT_FRAG_SAMPLE_COUNT, FragPopulationSummary, select_top_live_doc_candidates};

    #[test]
    fn select_top_live_doc_candidates_keeps_highest_live_doc_frags() {
        let mut candidates = vec![
            FragPopulationSummary {
                index: 0,
                estimated_live_docs: 10,
            },
            FragPopulationSummary {
                index: 1,
                estimated_live_docs: 90,
            },
            FragPopulationSummary {
                index: 2,
                estimated_live_docs: 20,
            },
            FragPopulationSummary {
                index: 3,
                estimated_live_docs: 80,
            },
            FragPopulationSummary {
                index: 4,
                estimated_live_docs: 30,
            },
            FragPopulationSummary {
                index: 5,
                estimated_live_docs: 70,
            },
            FragPopulationSummary {
                index: 6,
                estimated_live_docs: 40,
            },
            FragPopulationSummary {
                index: 7,
                estimated_live_docs: 60,
            },
            FragPopulationSummary {
                index: 8,
                estimated_live_docs: 50,
            },
            FragPopulationSummary {
                index: 9,
                estimated_live_docs: 100,
            },
        ];

        select_top_live_doc_candidates(&mut candidates, DEFAULT_FRAG_SAMPLE_COUNT);

        assert_eq!(candidates.len(), DEFAULT_FRAG_SAMPLE_COUNT);
        assert_eq!(
            candidates
                .into_iter()
                .map(|candidate| candidate.index)
                .collect::<Vec<_>>(),
            vec![9, 1, 3, 5, 7, 8, 6, 4]
        );
    }

    #[test]
    fn select_top_live_doc_candidates_breaks_ties_by_index() {
        let mut candidates = vec![
            FragPopulationSummary {
                index: 4,
                estimated_live_docs: 100,
            },
            FragPopulationSummary {
                index: 2,
                estimated_live_docs: 100,
            },
            FragPopulationSummary {
                index: 3,
                estimated_live_docs: 100,
            },
        ];

        select_top_live_doc_candidates(&mut candidates, DEFAULT_FRAG_SAMPLE_COUNT);

        assert_eq!(
            candidates
                .into_iter()
                .map(|candidate| candidate.index)
                .collect::<Vec<_>>(),
            vec![2, 3, 4]
        );
    }
}