use std::error::Error;

use log::{error, info, warn};
use tici_shard::reader::{
    query_info::QueryInfo,
    shard_mirror_manager::ShardMirrorManager,
    shard_reader::{
        EstimateSampleCounts, ShardEstimatePopulation, ShardEstimateSample, ShardReader,
    },
};

use crate::{
    build_query_info, build_read_ranges,
    ffi::{EstimateCountResult, SearchParam, ShardWithRange},
};

struct EstimateShardEntry {
    shard_id: u64,
    reader: ShardReader,
    query_info: QueryInfo,
    population: ShardEstimatePopulation,
}

#[derive(Default)]
struct AggregateEstimate {
    sampled_shards: u64,
    sampled_frags: u64,
    sampled_live_docs: u64,
    sampled_match_count: u64,
}

fn scale_by_ratio(base: u64, ratio_num: u64, ratio_den: u64) -> u64 {
    if base == 0 || ratio_num == 0 || ratio_den == 0 {
        0
    } else {
        (base as f64 / ratio_den as f64 * ratio_num as f64).round() as u64
    }
}

fn aggregate_sample(aggregate: &mut AggregateEstimate, sample: &ShardEstimateSample) {
    aggregate.sampled_shards += 1;
    aggregate.sampled_frags += sample.sampled_frags as u64;
    aggregate.sampled_live_docs += sample.sampled_live_docs;
    aggregate.sampled_match_count += sample.sampled_match_count;
}

fn build_result(
    estimated_total_count: u64,
    available_shard_count: u64,
    total_frags: u64,
    population_live_docs: u64,
    aggregate: &AggregateEstimate,
) -> EstimateCountResult {
    EstimateCountResult {
        estimated_total_count,
        available_shards: available_shard_count,
        sampled_shards: aggregate.sampled_shards,
        available_frags: total_frags,
        sampled_frags: aggregate.sampled_frags,
        population_live_docs,
        sampled_live_docs: aggregate.sampled_live_docs,
    }
}

fn select_sampled_shard_indexes(
    entries: &[EstimateShardEntry],
    shard_sample_count: usize,
) -> Vec<usize> {
    let mut candidates = entries
        .iter()
        .enumerate()
        .filter(|(_, entry)| entry.population.population_live_docs > 0)
        .map(|(index, entry)| (entry.population.population_live_docs, entry.shard_id, index))
        .collect::<Vec<_>>();
    candidates
        .sort_unstable_by(|left, right| right.0.cmp(&left.0).then_with(|| left.1.cmp(&right.1)));
    candidates.truncate(shard_sample_count);
    candidates
        .into_iter()
        .map(|(_, _, index)| index)
        .collect::<Vec<_>>()
}

pub fn estimate_count(
    shard_mirror_manager: &ShardMirrorManager,
    keyspace_id: u32,
    shards: &[ShardWithRange],
    query: &crate::ffi::Expr,
) -> Result<EstimateCountResult, Box<dyn Error>> {
    if shards.is_empty() {
        return Ok(EstimateCountResult::default());
    }

    let start = std::time::Instant::now();
    // Materialize a per-shard query context up front so the sampling stage can
    // reuse parsed expressions and the population summary for each shard.
    let count_search_param = SearchParam::default();
    let mut entries = Vec::with_capacity(shards.len());
    let mut total_frags = 0_u64;
    let mut population_live_docs = 0_u64;
    for shard in shards {
        let Some(shard_mirror) = shard_mirror_manager.get_shard(shard.shard_id) else {
            warn!("shard_id={} not found", shard.shard_id);
            continue;
        };
        let reader = shard_mirror.get_reader();
        let read_ranges = build_read_ranges(keyspace_id, &shard.ranges).inspect_err(|e| {
            error!(
                "Failed to create estimate read ranges, shard_id={}, err={:?}",
                shard.shard_id, e
            )
        })?;
        let query_info = build_query_info(&reader, query, &read_ranges, &[], &count_search_param)
            .inspect_err(|e| {
            error!(
                "Failed to parse estimate expr, shard_id={}, err={:?}",
                shard.shard_id, e
            )
        })?;
        let population = reader.estimate_population(query_info.read_ranges())?;
        total_frags += population.total_frags as u64;
        population_live_docs = population_live_docs.saturating_add(population.population_live_docs);
        entries.push(EstimateShardEntry {
            shard_id: shard.shard_id,
            reader,
            query_info,
            population,
        });
    }
    let available_shard_count = entries.len() as u64;

    let shard_ids = shards
        .iter()
        .map(|shard| shard.shard_id)
        .collect::<Vec<_>>();

    // If every shard reports zero live docs in the requested ranges, there is
    // nothing to sample or count. Returning an exact zero here avoids entering
    // the estimator with an empty population denominator.
    if population_live_docs == 0 {
        let result = build_result(
            0,
            available_shard_count,
            total_frags,
            0,
            &AggregateEstimate::default(),
        );
        info!(
            "Estimate keyspace_id={} success, shards={:?}, query={:?}, res={:?}, elapsed={}us",
            keyspace_id,
            shard_ids,
            query,
            result,
            start.elapsed().as_micros(),
        );
        return Ok(result);
    }

    // Sample up to `sample_counts.shard_sample_count` shards with the largest
    // live-doc populations in the requested ranges (see
    // `select_sampled_shard_indexes`). Within each sampled shard, the shard
    // reader evaluates the highest live-doc fragments first.
    let sample_counts = entries
        .first()
        .map(|entry| entry.reader.estimate_sample_counts())
        .unwrap_or_else(EstimateSampleCounts::default);
    let sampled_shard_indexes =
        select_sampled_shard_indexes(&entries, sample_counts.shard_sample_count);
    let mut aggregate = AggregateEstimate::default();
    for entry_index in sampled_shard_indexes {
        let sample = entries[entry_index]
            .reader
            .sample_count_with_query_info(&entries[entry_index].query_info)?;
        aggregate_sample(&mut aggregate, &sample);
    }
    let estimated_total_count = if aggregate.sampled_live_docs == 0 {
        0
    } else {
        scale_by_ratio(
            population_live_docs,
            aggregate.sampled_match_count,
            aggregate.sampled_live_docs,
        )
        .min(population_live_docs)
    };
    let result = build_result(
        estimated_total_count,
        available_shard_count,
        total_frags,
        population_live_docs,
        &aggregate,
    );
    info!(
        "Estimate keyspace_id={} success, shards={:?}, query={:?}, res={:?}, elapsed={}us",
        keyspace_id,
        shard_ids,
        query,
        result,
        start.elapsed().as_micros(),
    );
    Ok(result)
}

#[cfg(test)]
mod tests {
    use std::{fs, sync::Arc};

    use anyhow::{Context, Result, ensure};
    use tantivy::{Index, TantivyDocument};
    use tempfile::TempDir;
    use tici_common::codec::codec_v2::add_keyspace_prefix;
    use tici_grpc::{FragMeta, FragProperty, ShardManifestHeader};
    use tici_shard::{
        ShardMirrorMeta, ShardMirrorMetaVersion,
        config::FragReaderConfig,
        reader::{
            frag_reader::FragReaderMeta, shard_mirror_manager::ShardMirrorManager,
            shard_reader::DEFAULT_SHARD_SAMPLE_COUNT,
        },
        test_utils::{TEST_INDEX_KEY, TEST_INDEX_META_NORMAL, test_mock_frag_storage_refs_total},
        utils::{
            create_index_writer,
            dedup_bitset::FragAliveBitSetBytes,
            frag_index_path,
            schema::{HANDLE_STORE_NAME, VERSION_STORE_NAME},
        },
    };

    use super::*;

    struct TestShardArtifacts {
        _temp_dirs: Vec<TempDir>,
        shard_id: u64,
    }

    #[derive(Clone, Copy)]
    struct TestFragSpec {
        live_doc_count: usize,
        match_doc_count: usize,
    }

    fn new_manager(config: FragReaderConfig) -> Arc<ShardMirrorManager> {
        Arc::new(ShardMirrorManager::new(config))
    }

    fn default_manager() -> Arc<ShardMirrorManager> {
        new_manager(FragReaderConfig::default())
    }

    fn build_count_query_expr() -> crate::ffi::Expr {
        crate::ffi::Expr {
            tp: tipb::ExprType::ScalarFunc as i32,
            val: Vec::new(),
            children: vec![
                crate::ffi::Expr {
                    tp: tipb::ExprType::String as i32,
                    val: b"match".to_vec(),
                    children: Vec::new(),
                    rpn_args_len: 0,
                    sig: tipb::ScalarFuncSig::Unspecified as i32,
                    field_type: 0,
                },
                crate::ffi::Expr {
                    tp: tipb::ExprType::String as i32,
                    val: b"column_1".to_vec(),
                    children: Vec::new(),
                    rpn_args_len: 0,
                    sig: tipb::ScalarFuncSig::Unspecified as i32,
                    field_type: 0,
                },
            ],
            rpn_args_len: 0,
            sig: tipb::ScalarFuncSig::FtsMatchWord as i32,
            field_type: 0,
        }
    }

    fn build_missing_count_query_expr() -> crate::ffi::Expr {
        crate::ffi::Expr {
            tp: tipb::ExprType::ScalarFunc as i32,
            val: Vec::new(),
            children: vec![
                crate::ffi::Expr {
                    tp: tipb::ExprType::String as i32,
                    val: b"missing".to_vec(),
                    children: Vec::new(),
                    rpn_args_len: 0,
                    sig: tipb::ScalarFuncSig::Unspecified as i32,
                    field_type: 0,
                },
                crate::ffi::Expr {
                    tp: tipb::ExprType::String as i32,
                    val: b"column_1".to_vec(),
                    children: Vec::new(),
                    rpn_args_len: 0,
                    sig: tipb::ScalarFuncSig::Unspecified as i32,
                    field_type: 0,
                },
            ],
            rpn_args_len: 0,
            sig: tipb::ScalarFuncSig::FtsMatchWord as i32,
            field_type: 0,
        }
    }

    fn build_test_doc(handle: Vec<u8>, version: u64, text: &str) -> Result<TantivyDocument> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema();
        let mut doc = TantivyDocument::default();
        doc.add_bytes(schema.get_field(HANDLE_STORE_NAME)?, &handle);
        doc.add_u64(schema.get_field(VERSION_STORE_NAME)?, version);
        doc.add_text(schema.get_field("column_1")?, text);
        doc.add_text(schema.get_field("column_2")?, format!("desc_{version}"));
        doc.add_text(schema.get_field("column_3")?, format!("pk1_{version}"));
        doc.add_text(schema.get_field("column_4")?, format!("pk2_{version}"));
        index_meta.build_fulltext_fields(&mut doc)?;
        Ok(doc)
    }

    async fn add_uniform_test_shard(
        manager: Arc<ShardMirrorManager>,
        shard_id: u64,
        frag_count: usize,
    ) -> Result<TestShardArtifacts> {
        let frag_specs = vec![
            TestFragSpec {
                live_doc_count: 2,
                match_doc_count: 1,
            };
            frag_count
        ];
        add_test_shard(manager, shard_id, &frag_specs).await
    }

    async fn add_test_shard(
        manager: Arc<ShardMirrorManager>,
        shard_id: u64,
        frag_specs: &[TestFragSpec],
    ) -> Result<TestShardArtifacts> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let raw_start_key = format!("s{shard_id:02}_").into_bytes();
        let raw_end_key = format!("s{:02}_", shard_id + 1).into_bytes();
        let start_key = add_keyspace_prefix(TEST_INDEX_KEY.keyspace_id, &raw_start_key);
        let end_key = add_keyspace_prefix(TEST_INDEX_KEY.keyspace_id, &raw_end_key);

        let mut temp_dirs = Vec::with_capacity(frag_specs.len());
        let mut frag_reader_metas = Vec::with_capacity(frag_specs.len());
        for (frag_idx, frag_spec) in frag_specs.iter().copied().enumerate() {
            ensure!(frag_spec.live_doc_count > 0);
            ensure!(frag_spec.match_doc_count <= frag_spec.live_doc_count);

            let temp_dir = TempDir::new()?;
            let local_frag_path = temp_dir
                .path()
                .join(format!("frag_{shard_id}_{frag_idx}"))
                .to_string_lossy()
                .to_string();
            let index_path = frag_index_path(&local_frag_path);
            fs::create_dir_all(&index_path)
                .with_context(|| format!("create index dir failed: {index_path}"))?;

            let schema = index_meta.get_tantivy_schema().clone();
            let mut index = Index::create_in_dir(&index_path, schema)?;
            index.set_tokenizers(index_meta.get_analyzer_manager().clone());
            let mut writer = create_index_writer(&index, 1, 15_000_000, &local_frag_path)?;

            let mut handles = Vec::with_capacity(frag_spec.live_doc_count);
            for doc_idx in 0..frag_spec.live_doc_count {
                let raw_handle =
                    format!("s{shard_id:02}_f{frag_idx:02}_d{doc_idx:04}").into_bytes();
                let handle = add_keyspace_prefix(TEST_INDEX_KEY.keyspace_id, &raw_handle);
                let text = if doc_idx < frag_spec.match_doc_count {
                    "alpha match"
                } else {
                    "beta other"
                };
                writer.add_document(build_test_doc(
                    handle.clone(),
                    (frag_idx * 10_000 + doc_idx + 1) as u64,
                    text,
                )?)?;
                handles.push(handle);
            }
            writer.commit()?;
            writer.wait_merging_threads()?;

            let min_handle = handles[0].clone();
            let max_handle = handles[handles.len() - 1].clone();
            let split_handles = if handles.len() > 1 {
                vec![handles[handles.len() / 2].clone()]
            } else {
                vec![min_handle.clone()]
            };

            let frag_meta = FragMeta::new(
                1,
                format!("frag_{shard_id}_{frag_idx}/f_1"),
                FragProperty {
                    size: frag_spec.live_doc_count as u64,
                    count: frag_spec.live_doc_count as u64,
                    min_handle,
                    split_handles,
                    max_handle,
                    delete_count: 0,
                },
                test_mock_frag_storage_refs_total(frag_spec.live_doc_count as u64),
                1,
            )?;
            frag_reader_metas.push(FragReaderMeta::from_local_path(
                &frag_meta,
                local_frag_path,
            )?);
            temp_dirs.push(temp_dir);
        }

        manager
            .add_shard(
                ShardMirrorMeta {
                    version: ShardMirrorMetaVersion::LATEST,
                    index_key: TEST_INDEX_KEY.clone(),
                    shard_header: ShardManifestHeader {
                        shard_id,
                        start_key,
                        end_key,
                        epoch: 1,
                        seq: 1,
                    },
                    frag_bucket: "test".to_string(),
                },
                index_meta,
                frag_reader_metas,
                (0..frag_specs.len())
                    .map(|_| FragAliveBitSetBytes::default())
                    .collect(),
            )
            .await?;

        Ok(TestShardArtifacts {
            _temp_dirs: temp_dirs,
            shard_id,
        })
    }

    fn build_test_shards_with_ranges<'a>(
        shard_ids: &[u64],
        ranges: &[crate::ffi::Range<'a>],
    ) -> Vec<crate::ffi::ShardWithRange<'a>> {
        shard_ids
            .iter()
            .copied()
            .map(|shard_id| crate::ffi::ShardWithRange {
                shard_id,
                ranges: ranges.to_vec(),
            })
            .collect()
    }

    #[tokio::test]
    async fn test_estimate_count_small_scope_uses_sampling_result() -> Result<()> {
        let manager = default_manager();
        let fixtures = [
            add_uniform_test_shard(manager.clone(), 0, 2).await?,
            add_uniform_test_shard(manager.clone(), 1, 2).await?,
        ];
        let shard_ids = fixtures
            .iter()
            .map(|fixture| fixture.shard_id)
            .collect::<Vec<_>>();
        let ranges = vec![crate::ffi::Range {
            start: b"s00_",
            end: b"s09_",
        }];
        let shards = build_test_shards_with_ranges(&shard_ids, &ranges);

        let result = estimate_count(
            &manager,
            TEST_INDEX_KEY.keyspace_id,
            &shards,
            &build_count_query_expr(),
        )
        .map_err(|err| anyhow::anyhow!("{err}"))?;

        assert_eq!(result.estimated_total_count, 4);
        assert_eq!(result.available_shards, 2);
        assert_eq!(result.sampled_shards, 2);
        assert_eq!(result.available_frags, 4);
        assert_eq!(result.sampled_frags, 4);
        Ok(())
    }

    #[tokio::test]
    async fn test_estimate_count_sampling_path_returns_uniform_estimate() -> Result<()> {
        let manager = default_manager();
        let mut fixtures = Vec::new();
        for shard_id in 0..8_u64 {
            fixtures.push(add_uniform_test_shard(manager.clone(), shard_id, 4).await?);
        }
        let shard_ids = fixtures
            .iter()
            .map(|fixture| fixture.shard_id)
            .collect::<Vec<_>>();
        let ranges = vec![crate::ffi::Range {
            start: b"s00_",
            end: b"s99_",
        }];
        let shards = build_test_shards_with_ranges(&shard_ids, &ranges);

        let result = estimate_count(
            &manager,
            TEST_INDEX_KEY.keyspace_id,
            &shards,
            &build_count_query_expr(),
        )
        .map_err(|err| anyhow::anyhow!("{err}"))?;

        assert_eq!(result.estimated_total_count, 8 * 4);
        assert_eq!(result.available_shards, 8);
        assert_eq!(result.sampled_shards, DEFAULT_SHARD_SAMPLE_COUNT as u64);
        assert_eq!(
            result.sampled_frags,
            (DEFAULT_SHARD_SAMPLE_COUNT * 4) as u64
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_estimate_count_uses_custom_sample_counts() -> Result<()> {
        let manager = new_manager(FragReaderConfig {
            shard_sample_count: 2,
            frag_sample_count: 1,
            ..FragReaderConfig::default()
        });
        let mut fixtures = Vec::new();
        for shard_id in 0..4_u64 {
            let frag_specs = vec![
                TestFragSpec {
                    live_doc_count: 2,
                    match_doc_count: 1,
                };
                3
            ];
            fixtures.push(add_test_shard(manager.clone(), shard_id, &frag_specs).await?);
        }
        let shard_ids = fixtures
            .iter()
            .map(|fixture| fixture.shard_id)
            .collect::<Vec<_>>();
        let ranges = vec![crate::ffi::Range {
            start: b"s00_",
            end: b"s99_",
        }];
        let shards = build_test_shards_with_ranges(&shard_ids, &ranges);

        let result = estimate_count(
            &manager,
            TEST_INDEX_KEY.keyspace_id,
            &shards,
            &build_count_query_expr(),
        )
        .map_err(|err| anyhow::anyhow!("{err}"))?;

        assert_eq!(result.estimated_total_count, 12);
        assert_eq!(result.sampled_shards, 2);
        assert_eq!(result.sampled_frags, 2);
        Ok(())
    }

    #[tokio::test]
    async fn test_estimate_count_zero_hit_sampling_returns_zero() -> Result<()> {
        let manager = default_manager();
        let mut fixtures = Vec::new();
        for shard_id in 0..8_u64 {
            fixtures.push(add_uniform_test_shard(manager.clone(), shard_id, 4).await?);
        }
        let shard_ids = fixtures
            .iter()
            .map(|fixture| fixture.shard_id)
            .collect::<Vec<_>>();
        let ranges = vec![crate::ffi::Range {
            start: b"s00_",
            end: b"s99_",
        }];
        let shards = build_test_shards_with_ranges(&shard_ids, &ranges);

        let result = estimate_count(
            &manager,
            TEST_INDEX_KEY.keyspace_id,
            &shards,
            &build_missing_count_query_expr(),
        )
        .map_err(|err| anyhow::anyhow!("{err}"))?;

        assert_eq!(result.estimated_total_count, 0);
        assert_eq!(result.sampled_shards, DEFAULT_SHARD_SAMPLE_COUNT as u64);
        assert_eq!(
            result.sampled_frags,
            (DEFAULT_SHARD_SAMPLE_COUNT * 4) as u64
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_estimate_count_non_uniform_sampling_scales_top_live_doc_shards() -> Result<()> {
        let manager = default_manager();
        let shard_specs = [
            vec![
                TestFragSpec {
                    live_doc_count: 40,
                    match_doc_count: 20,
                },
                TestFragSpec {
                    live_doc_count: 20,
                    match_doc_count: 10,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 25,
                    match_doc_count: 25,
                },
                TestFragSpec {
                    live_doc_count: 25,
                    match_doc_count: 0,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 30,
                    match_doc_count: 0,
                },
                TestFragSpec {
                    live_doc_count: 15,
                    match_doc_count: 9,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 20,
                    match_doc_count: 10,
                },
                TestFragSpec {
                    live_doc_count: 20,
                    match_doc_count: 10,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 15,
                    match_doc_count: 9,
                },
                TestFragSpec {
                    live_doc_count: 20,
                    match_doc_count: 5,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 18,
                    match_doc_count: 3,
                },
                TestFragSpec {
                    live_doc_count: 12,
                    match_doc_count: 3,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 10,
                    match_doc_count: 10,
                },
                TestFragSpec {
                    live_doc_count: 10,
                    match_doc_count: 10,
                },
            ],
            vec![
                TestFragSpec {
                    live_doc_count: 10,
                    match_doc_count: 0,
                },
                TestFragSpec {
                    live_doc_count: 5,
                    match_doc_count: 0,
                },
            ],
        ];

        let mut fixtures = Vec::with_capacity(shard_specs.len());
        for (shard_id, frag_specs) in shard_specs.iter().enumerate() {
            fixtures.push(add_test_shard(manager.clone(), shard_id as u64, frag_specs).await?);
        }
        let shard_ids = fixtures
            .iter()
            .map(|fixture| fixture.shard_id)
            .collect::<Vec<_>>();
        let ranges = vec![crate::ffi::Range {
            start: b"s00_",
            end: b"s99_",
        }];
        let shards = build_test_shards_with_ranges(&shard_ids, &ranges);

        let result = estimate_count(
            &manager,
            TEST_INDEX_KEY.keyspace_id,
            &shards,
            &build_count_query_expr(),
        )
        .map_err(|err| anyhow::anyhow!("{err}"))?;

        assert_eq!(result.available_shards, shard_specs.len() as u64);
        assert_eq!(result.sampled_shards, DEFAULT_SHARD_SAMPLE_COUNT as u64);
        assert_eq!(result.population_live_docs, 295);
        assert_eq!(result.sampled_live_docs, 260);
        assert_eq!(
            result.sampled_frags,
            (DEFAULT_SHARD_SAMPLE_COUNT * 2) as u64
        );
        assert_eq!(result.estimated_total_count, 118);
        Ok(())
    }

    #[test]
    fn test_estimate_count_skips_missing_shard() {
        let manager = default_manager();
        let shard_ids = vec![42];
        let ranges = vec![crate::ffi::Range {
            start: b"",
            end: b"\xff",
        }];
        let shards = build_test_shards_with_ranges(&shard_ids, &ranges);

        let result = estimate_count(
            &manager,
            TEST_INDEX_KEY.keyspace_id,
            &shards,
            &build_count_query_expr(),
        )
        .unwrap_or_else(|err| panic!("missing shard should be skipped, got error: {err}"));
        assert_eq!(result.available_shards, 0);
        assert_eq!(result.sampled_shards, 0);
        assert_eq!(result.estimated_total_count, 0);
    }

    #[test]
    fn test_build_read_ranges_uses_test_index_key_keyspace() {
        let ranges = vec![crate::ffi::Range {
            start: b"",
            end: b"\xff",
        }];

        let built_ranges = build_read_ranges(TEST_INDEX_KEY.keyspace_id, &ranges)
            .unwrap_or_else(|err| panic!("build ranges should succeed, got error: {err}"));
        assert_eq!(built_ranges.len(), 1);
    }
}