#![allow(dead_code)]

use std::time::Duration;

use criterion::BenchmarkGroup;
use criterion::measurement::WallTime;
use rand::{RngExt, SeedableRng};

/// Change this alias to switch the value type used by every benchmark.
pub type ValueType = i64;

use qdrant_filter_test::inline;
use qdrant_filter_test::origin;

// ── Shared constants ─────────────────────────────────────────────────────────
// Change these once and every benchmark will pick up the new values.

/// Seed used to generate the point data.
pub const DATA_SEED: u64 = 42;

/// Seed used to generate random query point IDs.
pub const QUERY_SEED: u64 = 123;

/// Total number of points to generate.
pub const POINTS_NUM: usize = 4000_0000;

/// Upper bound (inclusive) for the random value assigned to each point.
pub const MAX_VALUE: ValueType = 2;

/// Number of point IDs to query during the benchmark.
pub const QUERY_COUNT: usize = 10_0000;

/// The value we pass to `check_values_any` as the predicate target.
pub const CHECK_VALUE: ValueType = 1;

/// Ratio of inner `Vec`s that are empty in the sparse-single-value data type.
/// Must be in the range `0.0..=1.0`. For example, `0.3` means 30 % of entries
/// will be empty.
pub const EMPTY_RATIO: f64 = 0.1;

/// Maximum number of elements in each inner `Vec` for the multi-value data
/// type. Each inner `Vec` will have a random length in `1..=MAX_VALUES_PER_POINT`.
pub const MAX_VALUES_PER_POINT: usize = 5;

// ── Data builders ────────────────────────────────────────────────────────────

/// **Type 1 – single value**: every inner `Vec` has exactly one element.
pub fn build_single_value_data(
    seed: u64,
    points_num: usize,
    max_value: ValueType,
) -> Vec<Vec<ValueType>> {
    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    let mut values = Vec::with_capacity(points_num);
    for _ in 0..points_num {
        let value: ValueType = rng.random_range(0..=max_value);
        values.push(vec![value]);
    }
    values
}

/// **Type 2 – sparse single value**: every inner `Vec` has *at most* one
/// element. A proportion of `empty_ratio` entries will be empty.
pub fn build_sparse_single_value_data(
    seed: u64,
    points_num: usize,
    max_value: ValueType,
    empty_ratio: f64,
) -> Vec<Vec<ValueType>> {
    use rand::seq::SliceRandom;

    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    let empty_count = (points_num as f64 * empty_ratio).round() as usize;
    let non_empty_count = points_num - empty_count;

    let mut values = Vec::with_capacity(points_num);

    // First, push all non-empty entries
    for _ in 0..non_empty_count {
        let value: ValueType = rng.random_range(0..=max_value);
        values.push(vec![value]);
    }

    // Then, push all empty entries
    for _ in 0..empty_count {
        values.push(vec![]);
    }

    // Shuffle to distribute empties randomly
    values.shuffle(&mut rng);

    values
}

/// **Type 3 – multi value**: every inner `Vec` has a random number of elements
/// in `1..=max_values_per_point`.
pub fn build_multi_value_data(
    seed: u64,
    points_num: usize,
    max_value: ValueType,
    max_values_per_point: usize,
) -> Vec<Vec<ValueType>> {
    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    let mut values = Vec::with_capacity(points_num);
    for _ in 0..points_num {
        let count = rng.random_range(1..=max_values_per_point);
        let inner: Vec<ValueType> = (0..count)
            .map(|_| rng.random_range(0..=max_value))
            .collect();
        values.push(inner);
    }
    values
}

// ── TestData ─────────────────────────────────────────────────────────────────

/// Pre-built instances of every implementation, ready for benchmarking.
pub struct TestData {
    pub inline: inline::ImmutablePointToValues<ValueType>,
    pub origin: origin::ImmutablePointToValues<ValueType>,
}

impl TestData {
    /// Construct all implementations from the given raw data.
    pub fn from_data(points_values: Vec<Vec<ValueType>>) -> Self {
        let inline = inline::ImmutablePointToValues::new(points_values.clone());
        let origin = origin::ImmutablePointToValues::new(points_values.clone());

        Self { inline, origin }
    }

    /// Convenience: build from the *single-value* data type using the shared
    /// constants.
    pub fn single_value() -> Self {
        Self::from_data(build_single_value_data(DATA_SEED, POINTS_NUM, MAX_VALUE))
    }

    /// Convenience: build from the *sparse-single-value* data type using the
    /// shared constants.
    pub fn sparse_single_value() -> Self {
        Self::from_data(build_sparse_single_value_data(
            DATA_SEED,
            POINTS_NUM,
            MAX_VALUE,
            EMPTY_RATIO,
        ))
    }

    /// Convenience: build from the *multi-value* data type using the shared
    /// constants.
    pub fn multi_value() -> Self {
        Self::from_data(build_multi_value_data(
            DATA_SEED,
            POINTS_NUM,
            MAX_VALUE,
            MAX_VALUES_PER_POINT,
        ))
    }
}

// ── Query-ID helpers ─────────────────────────────────────────────────────────

/// Generate `QUERY_COUNT` random point IDs in `[0, POINTS_NUM)`.
pub fn build_random_point_ids() -> Vec<u32> {
    let mut rng = rand::rngs::StdRng::seed_from_u64(QUERY_SEED);
    (0..QUERY_COUNT)
        .map(|_| rng.random_range(0..POINTS_NUM as u32))
        .collect()
}

// ── Bench Group Config ───────────────────────────────────────────────────────

/// Configure a benchmark group with standardised timing parameters.
pub fn configure_benchmark(group: &mut BenchmarkGroup<WallTime>) {
    group.measurement_time(Duration::from_secs(30));
}

mod common;

use criterion::{Criterion, black_box, criterion_group, criterion_main, configure_benchmark};

use crate::common::{CHECK_VALUE, TestData, build_random_point_ids};

// ── Helper: run all implementations against one TestData ─────────────────────

fn bench_all_impls(
    group: &mut criterion::BenchmarkGroup<criterion::measurement::WallTime>,
    data: &TestData,
    point_ids: &[u32],
) {
    configure_benchmark(group);

    group.bench_function("origin", |b| {
        b.iter(|| {
            for &point_id in point_ids {
                black_box(
                    data.origin
                        .check_values_any(black_box(point_id), |v| *v == CHECK_VALUE),
                );
            }
        });
    });

    group.bench_function("inline", |b| {
        b.iter(|| {
            for &point_id in point_ids {
                black_box(
                    data.inline
                        .check_values_any(black_box(point_id), |v| *v == CHECK_VALUE),
                );
            }
        });
    });
}

// ── Benchmarks ───────────────────────────────────────────────────────────────

fn random_single_value(c: &mut Criterion) {
    let data = TestData::single_value();
    let point_ids = build_random_point_ids();

    let mut group = c.benchmark_group("Random single_value");
    bench_all_impls(&mut group, &data, &point_ids);
    group.finish();
}

fn random_sparse_single_value(c: &mut Criterion) {
    let data = TestData::sparse_single_value();
    let point_ids = build_random_point_ids();

    let mut group = c.benchmark_group("Random sparse_single_value");
    bench_all_impls(&mut group, &data, &point_ids);
    group.finish();
}

fn random_multi_value(c: &mut Criterion) {
    let data = TestData::multi_value();
    let point_ids = build_random_point_ids();

    let mut group = c.benchmark_group("Random multi_value");
    bench_all_impls(&mut group, &data, &point_ids);
    group.finish();
}

criterion_group!(
    benches,
    random_single_value,
    random_sparse_single_value,
    random_multi_value,
);
criterion_main!(benches);

use std::ops::Range;

pub struct ImmutablePointToValues<N: Default> {
    point_to_values: Vec<Range<u32>>,
    point_to_values_container: Vec<N>,
}

impl<N: Default> ImmutablePointToValues<N> {
    pub fn new(src: Vec<Vec<N>>) -> Self {
        let mut point_to_values = Vec::with_capacity(src.len());
        let all_values_count = src.iter().fold(0, |acc, values| acc + values.len());
        let mut point_to_values_container = Vec::with_capacity(all_values_count);
        for values in src {
            let container_len = point_to_values_container.len() as u32;
            let range = container_len..container_len + values.len() as u32;
            point_to_values.push(range.clone());
            point_to_values_container.extend(values);
        }
        Self {
            point_to_values,
            point_to_values_container,
        }
    }

    pub fn check_values_any(&self, idx: u32, check_fn: impl FnMut(&N) -> bool) -> bool {
        let Some(range) = self.point_to_values.get(idx as usize) else {
            return false;
        };

        let range = range.start as usize..range.end as usize;
        if let Some(values) = self.point_to_values_container.get(range) {
            values.iter().any(check_fn)
        } else {
            false
        }
    }
}

#[derive(Debug, Clone)]
enum PointValueEntry<N> {
    Single(N),
    Slice { start: u32, count: u32 },
}

impl<N> Default for PointValueEntry<N> {
    fn default() -> Self {
        PointValueEntry::Slice { start: 0, count: 0 }
    }
}

#[derive(Debug, Clone, Default)]
pub struct ImmutablePointToValues<N: Default> {
    point_entries: Vec<PointValueEntry<N>>,
    values_container: Vec<N>,
}

impl<N: Default> ImmutablePointToValues<N> {
    pub fn new(src: Vec<Vec<N>>) -> Self {
        let mut point_entries = Vec::with_capacity(src.len());
        let container_capacity = src
            .iter()
            .filter(|values| values.len() != 1)
            .fold(0, |acc, values| acc + values.len());
        let mut values_container = Vec::with_capacity(container_capacity);

        for values in src {
            if values.len() == 1 {
                let value = values.into_iter().next().expect("length checked above");
                point_entries.push(PointValueEntry::Single(value));
            } else {
                let start = values_container.len() as u32;
                let count = values.len() as u32;
                point_entries.push(PointValueEntry::Slice { start, count });
                values_container.extend(values);
            }
        }

        Self {
            point_entries,
            values_container,
        }
    }

    pub fn check_values_any(&self, idx: u32, check_fn: impl FnMut(&N) -> bool) -> bool {
        let Some(entry) = self.point_entries.get(idx as usize) else {
            return false;
        };

        match entry {
            PointValueEntry::Single(v) => std::iter::once(v).any(check_fn),
            PointValueEntry::Slice { start, count } => {
                let range = *start as usize..(*start + *count) as usize;
                if let Some(values) = self.values_container.get(range) {
                    values.iter().any(check_fn)
                } else {
                    false
                }
            }
        }
    }
}