// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Context for iterating bulk memtable.

use std::collections::VecDeque;
use std::sync::Arc;

use common_recordbatch::filter::SimpleFilterEvaluator;
use mito_codec::row_converter::build_primary_key_codec;
use parquet::file::metadata::ParquetMetaData;
use store_api::metadata::RegionMetadataRef;
use store_api::storage::ColumnId;
use table::predicate::Predicate;

use crate::error::Result;
use crate::read::read_columns::ReadColumns;
use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
use crate::sst::parquet::file_range::{PreFilterMode, RangeBase};
use crate::sst::parquet::flat_format::FlatReadFormat;
use crate::sst::parquet::prefilter::{CachedPrimaryKeyFilter, build_bulk_filter_plan};
use crate::sst::parquet::stats::RowGroupPruningStats;

pub(crate) type BulkIterContextRef = Arc<BulkIterContext>;

pub struct BulkIterContext {
    pub(crate) base: RangeBase,
    pub(crate) predicate: Option<Predicate>,
    /// Pre-extracted primary key filters for PK prefiltering.
    /// `None` if PK prefiltering is not applicable.
    pk_filters: Option<Arc<Vec<SimpleFilterEvaluator>>>,
    batch_size: usize,
}

impl BulkIterContext {
    pub fn new(
        region_metadata: RegionMetadataRef,
        projection: Option<&[ColumnId]>,
        predicate: Option<Predicate>,
        skip_auto_convert: bool,
        batch_size: usize,
    ) -> Result<Self> {
        Self::new_with_pre_filter_mode(
            region_metadata,
            projection,
            predicate,
            skip_auto_convert,
            PreFilterMode::All,
            batch_size,
        )
    }

    pub fn new_with_pre_filter_mode(
        region_metadata: RegionMetadataRef,
        projection: Option<&[ColumnId]>,
        predicate: Option<Predicate>,
        skip_auto_convert: bool,
        pre_filter_mode: PreFilterMode,
        batch_size: usize,
    ) -> Result<Self> {
        let codec = build_primary_key_codec(&region_metadata);

        let read_cols = if let Some(col_ids) = projection {
            ReadColumns::new(col_ids.iter().copied())
        } else {
            ReadColumns::new(
                region_metadata
                    .column_metadatas
                    .iter()
                    .map(|col| col.column_id),
            )
        };
        let read_format = FlatReadFormat::new(
            region_metadata.clone(),
            read_cols,
            None,
            "memtable",
            skip_auto_convert,
        )?;

        let dyn_filters = predicate
            .as_ref()
            .map(|pred| pred.dyn_filters().as_ref().clone())
            .unwrap_or_default();

        let filter_plan = build_bulk_filter_plan(&read_format, predicate.as_ref());

        Ok(Self {
            base: RangeBase {
                filters: filter_plan.remaining_simple_filters,
                dyn_filters,
                read_format,
                prune_schema: region_metadata.schema.clone(),
                expected_metadata: Some(region_metadata),
                codec,
                // we don't need to compat batch since all batch in memtable have the same schema.
                compat_batch: None,
                compaction_projection_mapper: None,
                pre_filter_mode,
                partition_filter: None,
            },
            predicate,
            pk_filters: filter_plan.pk_filters,
            batch_size: batch_size.clamp(1, DEFAULT_READ_BATCH_SIZE),
        })
    }

    pub(crate) fn batch_size(&self) -> usize {
        self.batch_size
    }

    /// Prunes row groups by stats.
    pub(crate) fn row_groups_to_read(
        &self,
        file_meta: &Arc<ParquetMetaData>,
        skip_fields: bool,
    ) -> VecDeque<usize> {
        let region_meta = self.base.read_format.metadata();
        let row_groups = file_meta.row_groups();
        // expected_metadata is set to None since we always expect region metadata of memtable is up-to-date.
        let stats =
            RowGroupPruningStats::new(row_groups, &self.base.read_format, None, skip_fields);
        if let Some(predicate) = self.predicate.as_ref() {
            predicate
                .prune_with_stats(&stats, region_meta.schema.arrow_schema())
                .iter()
                .zip(0..file_meta.num_row_groups())
                .filter_map(|(selected, row_group)| {
                    if !*selected {
                        return None;
                    }
                    Some(row_group)
                })
                .collect::<VecDeque<_>>()
        } else {
            (0..file_meta.num_row_groups()).collect()
        }
    }

    /// Builds a fresh PK filter for a new iterator. Returns `None` if PK
    /// prefiltering is not applicable.
    pub(crate) fn build_pk_filter(&self) -> Option<CachedPrimaryKeyFilter> {
        let pk_filters = self.pk_filters.as_ref()?;
        let metadata = self.base.read_format.metadata();
        let inner = self
            .base
            .codec
            .primary_key_filter(metadata, Arc::clone(pk_filters));
        Some(CachedPrimaryKeyFilter::new(inner))
    }

    pub(crate) fn read_format(&self) -> &FlatReadFormat {
        &self.base.read_format
    }

    /// Returns the pre-filter mode.
    pub(crate) fn pre_filter_mode(&self) -> PreFilterMode {
        self.base.pre_filter_mode
    }

    /// Returns the region id.
    pub(crate) fn region_id(&self) -> store_api::storage::RegionId {
        self.base.read_format.metadata().region_id
    }
}
