// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

use crate::distance::{preprocess_vectors, MetricType};
use crate::kmeans::KMeansConfig;
use crate::pq::ProductQuantizer;
use crate::vamana::{
    estimate_sharded_vamana_memory_bytes, estimate_vamana_memory_bytes, VamanaGraph,
};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use std::borrow::Cow;
use std::collections::VecDeque;
use std::io;
use std::time::{Duration, Instant};

pub(crate) const DISKANN_ADJACENCY_LOCATOR_NODE_BYTES: usize = 4;
pub(crate) const DISKANN_ADJACENCY_LOCATOR_BLOCK_NODES: usize = 16;
/// Match the proven DiskANN training bound: more samples materially increase
/// memory and training time without consistently improving the codebook.
pub const DISKANN_MAX_PQ_TRAINING_VECTORS: usize = 50_000;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DiskAnnStorageLayout {
    /// Keep compressed adjacency pages and dense raw-vector records in separate sections.
    Compact,
    /// Store each raw vector immediately before its compressed adjacency list.
    Interleaved,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
pub enum DiskAnnRawVectorEncoding {
    /// Preserve indexed vectors and final distances as little-endian `f32`.
    F32 = 1,
    /// Store little-endian IEEE 754 binary16 values for approximate final reranking.
    F16 = 2,
}

impl DiskAnnRawVectorEncoding {
    pub(crate) const fn element_size(self) -> usize {
        match self {
            Self::F32 => size_of::<f32>(),
            Self::F16 => size_of::<u16>(),
        }
    }

    pub(crate) const fn from_code(code: u32) -> Option<Self> {
        match code {
            1 => Some(Self::F32),
            2 => Some(Self::F16),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DiskAnnBuildDistance {
    /// Use full-precision distances for graph traversal and robust pruning.
    FullPrecision,
    /// Use PQ distances for graph traversal and full precision for robust pruning.
    ProductQuantized,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct DiskAnnBuildParams {
    pub max_degree: usize,
    pub build_search_list_size: usize,
    pub alpha: f32,
    pub seed: u64,
    pub memory_budget_bytes: usize,
    pub storage_layout: DiskAnnStorageLayout,
    pub raw_vector_encoding: DiskAnnRawVectorEncoding,
    pub build_distance: DiskAnnBuildDistance,
}

impl Default for DiskAnnBuildParams {
    fn default() -> Self {
        Self {
            max_degree: 64,
            build_search_list_size: 100,
            alpha: 1.2,
            seed: 42,
            memory_budget_bytes: 8 * 1024 * 1024 * 1024,
            storage_layout: DiskAnnStorageLayout::Compact,
            raw_vector_encoding: DiskAnnRawVectorEncoding::F16,
            build_distance: DiskAnnBuildDistance::ProductQuantized,
        }
    }
}

pub(crate) fn validate_diskann_format_configuration(
    dimension: usize,
    pq_m: usize,
    pq_bits: usize,
    build: DiskAnnBuildParams,
) -> io::Result<()> {
    if dimension == 0 {
        return Err(invalid_input("DiskANN dimension must be greater than 0"));
    }
    if dimension > 1024 {
        return Err(invalid_input("DiskANN v1 dimension must be at most 1024"));
    }
    if pq_m == 0 {
        return Err(invalid_input("DiskANN pq.m must be greater than 0"));
    }
    if pq_m > dimension {
        return Err(invalid_input(format!(
            "DiskANN pq.m {} must not exceed dimension {}",
            pq_m, dimension
        )));
    }
    if !matches!(pq_bits, 4 | 8) {
        return Err(invalid_input("DiskANN pq.bits must be 4 or 8"));
    }
    if build.max_degree == 0 {
        return Err(invalid_input(
            "DiskANN maximum degree must be greater than 0",
        ));
    }
    if build.max_degree > 1023 {
        return Err(invalid_input(format!(
            "DiskANN adjacency list size {} exceeds the v1 1023-neighbor page limit",
            build.max_degree.saturating_mul(size_of::<u32>())
        )));
    }
    if build.build_search_list_size < build.max_degree {
        return Err(invalid_input(format!(
            "DiskANN build search-list size {} must be at least maximum degree {}",
            build.build_search_list_size, build.max_degree
        )));
    }
    if u32::try_from(build.build_search_list_size).is_err() {
        return Err(invalid_input("DiskANN build search-list size exceeds u32"));
    }
    if !build.alpha.is_finite() || build.alpha < 1.0 {
        return Err(invalid_input("DiskANN alpha must be at least 1 and finite"));
    }
    let interleaved_record_bytes = dimension
        .checked_mul(build.raw_vector_encoding.element_size())
        .and_then(|vector_bytes| {
            build
                .max_degree
                .checked_mul(size_of::<u32>())
                .and_then(|adjacency_bytes| vector_bytes.checked_add(adjacency_bytes))
        });
    if build.storage_layout == DiskAnnStorageLayout::Interleaved
        && interleaved_record_bytes.is_none_or(|record_bytes| record_bytes > 4096)
    {
        return Err(invalid_input(
            "DiskANN interleaved raw vector and maximum adjacency list must fit in one page",
        ));
    }
    Ok(())
}

pub(crate) fn validate_diskann_training_budget(
    dimension: usize,
    metric: MetricType,
    pq_m: usize,
    pq_bits: usize,
    memory_budget_bytes: usize,
) -> io::Result<()> {
    let minimum_training_vectors = 1usize << pq_bits;
    pq_training_plan_with_sample_buffers(
        dimension,
        pq_m,
        minimum_training_vectors,
        minimum_training_vectors,
        memory_budget_bytes,
        usize::from(metric == MetricType::Cosine) + 1,
    )
    .map(|_| ())
    .map_err(|error| {
        invalid_input(format!(
            "DiskANN memory budget cannot fit minimum PQ training: {error}"
        ))
    })
}

pub(crate) fn diskann_training_sample_limit(
    dimension: usize,
    metric: MetricType,
    pq_m: usize,
    pq_bits: usize,
    memory_budget_bytes: usize,
) -> io::Result<usize> {
    pq_training_plan_with_sample_buffers(
        dimension,
        pq_m,
        1usize << pq_bits,
        DISKANN_MAX_PQ_TRAINING_VECTORS,
        memory_budget_bytes,
        usize::from(metric == MetricType::Cosine) + 1,
    )
    .map(|plan| plan.sample_count)
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct DiskAnnBuildStats {
    /// One for the normal parallel build; greater than one when the memory
    /// budget selected overlapping shard construction.
    pub graph_shards: usize,
    pub total: Duration,
    pub pq_encoding: Duration,
    pub vamana_initialization: Duration,
    pub vamana_pass_one: Duration,
    pub vamana_pass_two: Duration,
    pub connectivity_repair: Duration,
    pub locality_remap: Duration,
    pub resident_serialization: Duration,
    pub adjacency_serialization: Duration,
    pub vector_serialization: Duration,
}

impl DiskAnnBuildStats {
    pub fn accounted_duration(self) -> Duration {
        [
            self.pq_encoding,
            self.vamana_initialization,
            self.vamana_pass_one,
            self.vamana_pass_two,
            self.connectivity_repair,
            self.locality_remap,
            self.resident_serialization,
            self.adjacency_serialization,
            self.vector_serialization,
        ]
        .into_iter()
        .sum()
    }
}

pub struct DiskAnnIndex {
    pub d: usize,
    pub metric: MetricType,
    pub pq: ProductQuantizer,
    pub build_params: DiskAnnBuildParams,
    pub ids: Vec<i64>,
    pub vectors: Vec<f32>,
}

impl DiskAnnIndex {
    pub fn new(
        d: usize,
        metric: MetricType,
        pq_m: usize,
        build_params: DiskAnnBuildParams,
    ) -> Self {
        Self::with_pq_bits(d, metric, pq_m, 8, build_params)
    }

    pub fn with_pq_bits(
        d: usize,
        metric: MetricType,
        pq_m: usize,
        pq_bits: usize,
        build_params: DiskAnnBuildParams,
    ) -> Self {
        Self {
            d,
            metric,
            pq: ProductQuantizer::with_nbits_balanced(d, pq_m, pq_bits),
            build_params,
            ids: Vec::new(),
            vectors: Vec::new(),
        }
    }

    pub fn train(&mut self, data: &[f32], n: usize) -> io::Result<()> {
        if n == 0 {
            return Err(invalid_input(
                "DiskANN training vector count must be greater than zero",
            ));
        }
        let expected_len = n
            .checked_mul(self.d)
            .ok_or_else(|| invalid_input("DiskANN training data length overflows usize"))?;
        if data.len() != expected_len {
            return Err(invalid_input(format!(
                "DiskANN training data length {} does not match n * dimension {}",
                data.len(),
                expected_len
            )));
        }
        let plan = self.training_plan(n)?;
        let sample =
            bounded_pq_training_sample(data, n, self.d, plan.sample_count, self.build_params.seed);
        let sample = sample
            .as_deref()
            .unwrap_or(&data[..plan.sample_count.saturating_mul(self.d)]);
        let processed = self.preprocess_vectors(sample, plan.sample_count);
        self.pq.train_hot_start_with_parallelism(
            &processed,
            plan.sample_count,
            &KMeansConfig::default(),
            false,
            plan.parallelism,
        );
        Ok(())
    }

    fn training_plan(&self, n: usize) -> io::Result<PqTrainingPlan> {
        pq_training_plan_with_sample_buffers(
            self.d,
            self.pq.m,
            self.pq.ksub,
            n,
            self.build_params.memory_budget_bytes,
            usize::from(self.metric == MetricType::Cosine) + 1,
        )
    }

    pub fn add(&mut self, data: &[f32], ids: &[i64]) {
        self.ids.extend_from_slice(ids);
        self.vectors
            .extend_from_slice(self.preprocess_vectors(data, ids.len()).as_ref());
    }

    pub fn estimate_build_memory_bytes(&self) -> io::Result<usize> {
        let n = self.ids.len();
        let workers = rayon::current_num_threads().max(1);
        let raw_vectors = checked_bytes(self.vectors.len(), size_of::<f32>(), "raw vectors")?;
        let row_ids = checked_bytes(n, size_of::<i64>(), "row IDs")?;
        let row_id_encoding_scratch = row_id_encoding_scratch_bytes(n)?;
        let pq_codes = checked_bytes(n, self.pq.code_size(), "PQ codes")?;
        let pq_codebook = checked_bytes(self.pq.centroids.len(), size_of::<f32>(), "PQ codebook")?;
        let pq_build_distances = if self.build_params.build_distance
            == DiskAnnBuildDistance::ProductQuantized
        {
            self.pq
                .m
                .checked_mul(self.pq.ksub)
                .and_then(|value| value.checked_mul(self.pq.ksub))
                .and_then(|value| value.checked_mul(size_of::<f32>()))
                .ok_or_else(|| invalid_input("DiskANN PQ build-distance table size overflows"))?
        } else {
            0
        };
        let row_id_order = checked_bytes(n, size_of::<u32>(), "row-ID order")?;
        let adjacency_index = checked_bytes(
            n,
            DISKANN_ADJACENCY_LOCATOR_NODE_BYTES,
            "adjacency locators",
        )?
        .checked_add(checked_bytes(
            n.div_ceil(DISKANN_ADJACENCY_LOCATOR_BLOCK_NODES),
            size_of::<u64>(),
            "adjacency locator block offsets",
        )?)
        .ok_or_else(|| invalid_input("DiskANN adjacency index size overflows usize"))?;
        let vamana = estimate_vamana_memory_bytes(
            n,
            self.build_params.max_degree,
            self.build_params.build_search_list_size,
            workers,
        )
        .ok_or_else(|| invalid_input("DiskANN Vamana memory estimate overflows usize"))?;
        let graph_stage_peak = vamana.build_peak_bytes.max(vamana.remap_peak_bytes);

        [
            raw_vectors,
            row_ids,
            row_id_encoding_scratch,
            pq_codes,
            pq_codebook,
            pq_build_distances,
            row_id_order,
            adjacency_index,
            graph_stage_peak,
        ]
        .into_iter()
        .try_fold(0usize, |total, value| {
            total
                .checked_add(value)
                .ok_or_else(|| invalid_input("DiskANN memory estimate overflows usize"))
        })
    }

    fn graph_build_shard_count(&self) -> io::Result<usize> {
        let estimated = self.estimate_build_memory_bytes()?;
        if estimated <= self.build_params.memory_budget_bytes {
            return Ok(1);
        }
        let n = self.ids.len();
        let workers = rayon::current_num_threads().max(1);
        let vamana = estimate_vamana_memory_bytes(
            n,
            self.build_params.max_degree,
            self.build_params.build_search_list_size,
            workers,
        )
        .ok_or_else(|| invalid_input("DiskANN Vamana memory estimate overflows usize"))?;
        let graph_peak = vamana.build_peak_bytes.max(vamana.remap_peak_bytes);
        let fixed_bytes = estimated
            .checked_sub(graph_peak)
            .ok_or_else(|| invalid_input("DiskANN fixed memory estimate underflows"))?;
        let max_shards = 64.min(n / 2);
        for shard_count in 2..=max_shards {
            let Some(sharded_graph) = estimate_sharded_vamana_memory_bytes(
                n,
                self.d,
                self.build_params.max_degree.min(n.saturating_sub(1)),
                shard_count,
                if self.build_params.build_distance == DiskAnnBuildDistance::ProductQuantized {
                    self.pq.code_size()
                } else {
                    0
                },
            ) else {
                continue;
            };
            if fixed_bytes
                .checked_add(sharded_graph)
                .is_some_and(|peak| peak <= self.build_params.memory_budget_bytes)
            {
                return Ok(shard_count);
            }
        }
        Err(invalid_input(format!(
            "DiskANN estimated build memory {} exceeds memory budget {}; overlapping sharded build also cannot fit",
            estimated, self.build_params.memory_budget_bytes
        )))
    }

    pub(crate) fn validate_for_write(&self) -> io::Result<()> {
        validate_diskann_format_configuration(self.d, self.pq.m, self.pq.nbits, self.build_params)?;
        validate_diskann_training_budget(
            self.d,
            self.metric,
            self.pq.m,
            self.pq.nbits,
            self.build_params.memory_budget_bytes,
        )?;
        if self.build_params.memory_budget_bytes == 0 {
            return Err(invalid_input(
                "DiskANN memory budget must be greater than zero",
            ));
        }
        if self.ids.is_empty() || u32::try_from(self.ids.len()).is_err() {
            return Err(invalid_input(
                "DiskANN vector count must be between 1 and u32::MAX",
            ));
        }
        let expected_vectors = self
            .ids
            .len()
            .checked_mul(self.d)
            .ok_or_else(|| invalid_input("DiskANN vector shape overflows usize"))?;
        if self.vectors.len() != expected_vectors {
            return Err(invalid_input(format!(
                "DiskANN vector length {} does not match {} row IDs * dimension {}",
                self.vectors.len(),
                self.ids.len(),
                self.d
            )));
        }
        if let Some(offset) = self.vectors.iter().position(|value| !value.is_finite()) {
            return Err(invalid_input(format!(
                "DiskANN vector data contains a non-finite value at offset {}",
                offset
            )));
        }
        if self.build_params.raw_vector_encoding == DiskAnnRawVectorEncoding::F16 {
            if let Some(offset) = self
                .vectors
                .iter()
                .position(|&value| !half::f16::from_f32(value).is_finite())
            {
                return Err(invalid_input(format!(
                    "DiskANN vector data at offset {} is outside the finite f16 range",
                    offset
                )));
            }
        }

        let expected_ksub = 1usize
            .checked_shl(self.pq.nbits as u32)
            .ok_or_else(|| invalid_input("DiskANN PQ centroid count overflows usize"))?;
        let expected_centroids = self
            .d
            .checked_mul(expected_ksub)
            .ok_or_else(|| invalid_input("DiskANN PQ codebook shape overflows usize"))?;
        if self.pq.d != self.d
            || self.pq.ksub != expected_ksub
            || self.pq.centroids.len() != expected_centroids
            || !self.pq.has_valid_layout()
        {
            return Err(invalid_input("DiskANN PQ codebook shape is invalid"));
        }
        if let Some(offset) = self
            .pq
            .centroids
            .iter()
            .position(|value| !value.is_finite())
        {
            return Err(invalid_input(format!(
                "DiskANN PQ codebook contains a non-finite value at offset {}",
                offset
            )));
        }
        Ok(())
    }

    pub(crate) fn prepare_build(&self) -> io::Result<PreparedDiskAnn> {
        self.validate_for_write()?;
        let graph_shards = self.graph_build_shard_count()?;

        let pq_started = Instant::now();
        let mut pq_codes = vec![0u8; self.ids.len() * self.pq.code_size()];
        self.pq
            .encode_batch(&self.vectors, self.ids.len(), &mut pq_codes);
        let pq_encoding = pq_started.elapsed();
        let (mut graph, vamana_stats) = if graph_shards > 1 {
            match self.build_params.build_distance {
                DiskAnnBuildDistance::FullPrecision => VamanaGraph::build_sharded_with_stats(
                    &self.vectors,
                    self.ids.len(),
                    self.d,
                    self.graph_metric(),
                    self.build_params,
                    graph_shards,
                )?,
                DiskAnnBuildDistance::ProductQuantized => VamanaGraph::build_sharded_with_pq_stats(
                    &self.vectors,
                    &self.pq,
                    &pq_codes,
                    self.ids.len(),
                    self.d,
                    self.graph_metric(),
                    self.build_params,
                    graph_shards,
                )?,
            }
        } else {
            match self.build_params.build_distance {
                DiskAnnBuildDistance::FullPrecision => VamanaGraph::build_with_stats(
                    &self.vectors,
                    self.ids.len(),
                    self.d,
                    self.graph_metric(),
                    self.build_params,
                )?,
                DiskAnnBuildDistance::ProductQuantized => VamanaGraph::build_with_pq_stats(
                    &self.vectors,
                    &self.pq,
                    &pq_codes,
                    self.ids.len(),
                    self.d,
                    self.graph_metric(),
                    self.build_params,
                )?,
            }
        };
        let locality_started = Instant::now();
        let permutation = bfs_locality_permutation(&graph);
        remap_graph_in_place(&mut graph, &permutation);
        let locality_remap = locality_started.elapsed();
        Ok(PreparedDiskAnn {
            graph,
            permutation,
            pq_codes,
            stats: DiskAnnBuildStats {
                graph_shards,
                pq_encoding,
                vamana_initialization: vamana_stats.initialization,
                vamana_pass_one: vamana_stats.pass_one,
                vamana_pass_two: vamana_stats.pass_two,
                connectivity_repair: vamana_stats.connectivity_repair,
                locality_remap,
                ..DiskAnnBuildStats::default()
            },
        })
    }

    fn graph_metric(&self) -> MetricType {
        match self.metric {
            // Cosine vectors are normalized during training and ingestion, so
            // squared L2 has identical ordering and a well-behaved prune ratio.
            MetricType::Cosine => MetricType::L2,
            metric => metric,
        }
    }

    fn preprocess_vectors<'a>(&self, data: &'a [f32], count: usize) -> Cow<'a, [f32]> {
        if self.metric == MetricType::Cosine {
            Cow::Owned(preprocess_vectors(data, count, self.d, self.metric))
        } else {
            Cow::Borrowed(&data[..count * self.d])
        }
    }
}

fn bounded_pq_training_sample(
    data: &[f32],
    count: usize,
    dimension: usize,
    sample_limit: usize,
    seed: u64,
) -> Option<Vec<f32>> {
    if count <= sample_limit {
        return None;
    }
    let mut sample = Vec::with_capacity(sample_limit.saturating_mul(dimension));
    sample.extend_from_slice(&data[..sample_limit * dimension]);
    let mut rng = StdRng::seed_from_u64(seed);
    for (stream_index, vector) in data
        .chunks_exact(dimension)
        .enumerate()
        .skip(sample_limit)
        .take(count - sample_limit)
    {
        let replacement = rng.gen_range(0..=stream_index);
        if replacement < sample_limit {
            let start = replacement * dimension;
            sample[start..start + dimension].copy_from_slice(vector);
        }
    }
    Some(sample)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct PqTrainingPlan {
    sample_count: usize,
    parallelism: usize,
}

#[cfg(test)]
fn pq_training_plan(
    dimension: usize,
    pq_m: usize,
    ksub: usize,
    vector_count: usize,
    memory_budget_bytes: usize,
) -> io::Result<PqTrainingPlan> {
    pq_training_plan_with_sample_buffers(
        dimension,
        pq_m,
        ksub,
        vector_count,
        memory_budget_bytes,
        1,
    )
}

fn pq_training_plan_with_sample_buffers(
    dimension: usize,
    pq_m: usize,
    ksub: usize,
    vector_count: usize,
    memory_budget_bytes: usize,
    sample_buffer_count: usize,
) -> io::Result<PqTrainingPlan> {
    if vector_count == 0 {
        return Ok(PqTrainingPlan {
            sample_count: 0,
            parallelism: 1,
        });
    }
    let max_sample = vector_count.min(DISKANN_MAX_PQ_TRAINING_VECTORS);
    let min_sample = vector_count.min(ksub.max(1));
    let max_chunk_dimension = dimension.div_ceil(pq_m.max(1));
    let current_parallelism = rayon::current_num_threads().max(1).min(pq_m.max(1));
    let peak_for = |sample_count: usize, parallelism: usize| -> Option<usize> {
        let sample_bytes = sample_count
            .checked_mul(dimension)?
            .checked_mul(size_of::<f32>())?
            .checked_mul(sample_buffer_count)?;
        let codebook_bytes = dimension.checked_mul(ksub)?.checked_mul(size_of::<f32>())?;
        let subvector_copy_bytes = sample_count
            .checked_mul(max_chunk_dimension)?
            .checked_mul(2 * size_of::<f32>())?;
        let assignment_bytes = sample_count.checked_mul(size_of::<usize>() + size_of::<f32>())?;
        let score_matrix_bytes = sample_count
            .checked_mul(ksub)?
            .min(4 * 1024 * 1024)
            .checked_mul(size_of::<f32>())?;
        let centroid_scratch_bytes = max_chunk_dimension
            .checked_mul(ksub)?
            .checked_mul(3 * size_of::<f32>())?;
        let per_task = [
            subvector_copy_bytes,
            assignment_bytes,
            score_matrix_bytes,
            centroid_scratch_bytes,
        ]
        .into_iter()
        .try_fold(0usize, |total, bytes| total.checked_add(bytes))?;
        sample_bytes
            .checked_add(codebook_bytes)?
            .checked_add(per_task.checked_mul(parallelism)?)
    };
    if peak_for(min_sample, 1).is_none_or(|peak| peak > memory_budget_bytes) {
        return Err(invalid_input(format!(
            "{} bytes is below the estimated one-worker PQ training peak",
            memory_budget_bytes
        )));
    }

    let mut low = min_sample;
    let mut high = max_sample;
    while low < high {
        let middle = low + (high - low).div_ceil(2);
        if peak_for(middle, 1).is_some_and(|peak| peak <= memory_budget_bytes) {
            low = middle;
        } else {
            high = middle - 1;
        }
    }
    let sample_count = low;
    let parallelism = (1..=current_parallelism)
        .rev()
        .find(|&parallelism| {
            peak_for(sample_count, parallelism).is_some_and(|peak| peak <= memory_budget_bytes)
        })
        .unwrap_or(1);
    Ok(PqTrainingPlan {
        sample_count,
        parallelism,
    })
}

fn checked_bytes(count: usize, item_size: usize, name: &str) -> io::Result<usize> {
    count
        .checked_mul(item_size)
        .ok_or_else(|| invalid_input(format!("DiskANN {} byte size overflows usize", name)))
}

fn row_id_encoding_scratch_bytes(count: usize) -> io::Result<usize> {
    checked_bytes(count, size_of::<i64>(), "row-ID encoding scratch")
}

fn invalid_input(message: impl Into<String>) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidInput, message.into())
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct LocalityPermutation {
    pub(crate) new_to_old: Vec<u32>,
    pub(crate) old_to_new: Vec<u32>,
}

pub(crate) fn bfs_locality_permutation(graph: &VamanaGraph) -> LocalityPermutation {
    let count = graph.adjacency.len();
    let mut visited = vec![false; count];
    let mut new_to_old = Vec::with_capacity(count);
    let entry = graph.entry_node as usize;
    if entry < count {
        visited[entry] = true;
        let mut queue = VecDeque::from([entry]);
        while let Some(node) = queue.pop_front() {
            new_to_old.push(node as u32);
            for &neighbor in &graph.adjacency[node] {
                let neighbor = neighbor as usize;
                if neighbor < count && !visited[neighbor] {
                    visited[neighbor] = true;
                    queue.push_back(neighbor);
                }
            }
        }
    }
    for (node, was_visited) in visited.into_iter().enumerate() {
        if !was_visited {
            new_to_old.push(node as u32);
        }
    }

    let mut old_to_new = vec![0u32; count];
    for (new_id, &old_id) in new_to_old.iter().enumerate() {
        old_to_new[old_id as usize] = new_id as u32;
    }
    LocalityPermutation {
        new_to_old,
        old_to_new,
    }
}

pub(crate) fn remap_graph_in_place(graph: &mut VamanaGraph, permutation: &LocalityPermutation) {
    graph
        .adjacency
        .permute_and_map_neighbors(&permutation.old_to_new);
    graph.entry_node = permutation.old_to_new[graph.entry_node as usize];
}

pub(crate) struct PreparedDiskAnn {
    pub(crate) graph: VamanaGraph,
    pub(crate) permutation: LocalityPermutation,
    pub(crate) pq_codes: Vec<u8>,
    pub(crate) stats: DiskAnnBuildStats,
}

impl PreparedDiskAnn {
    pub(crate) fn row_id(&self, index: &DiskAnnIndex, new_id: usize) -> i64 {
        index.ids[self.permutation.new_to_old[new_id] as usize]
    }

    pub(crate) fn vector<'a>(&self, index: &'a DiskAnnIndex, new_id: usize) -> &'a [f32] {
        let old_id = self.permutation.new_to_old[new_id] as usize;
        &index.vectors[old_id * index.d..(old_id + 1) * index.d]
    }

    pub(crate) fn pq_code<'a>(&'a self, index: &DiskAnnIndex, new_id: usize) -> &'a [u8] {
        let old_id = self.permutation.new_to_old[new_id] as usize;
        let code_size = index.pq.code_size();
        &self.pq_codes[old_id * code_size..(old_id + 1) * code_size]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn low_level_diskann_training_uses_the_same_bounded_deterministic_reservoir() {
        let dimension = 2;
        let count = DISKANN_MAX_PQ_TRAINING_VECTORS + 257;
        let data = (0..count * dimension)
            .map(|offset| offset as f32)
            .collect::<Vec<_>>();

        let first = bounded_pq_training_sample(
            &data,
            count,
            dimension,
            DISKANN_MAX_PQ_TRAINING_VECTORS,
            73,
        )
        .unwrap();
        let second = bounded_pq_training_sample(
            &data,
            count,
            dimension,
            DISKANN_MAX_PQ_TRAINING_VECTORS,
            73,
        )
        .unwrap();
        let different_seed = bounded_pq_training_sample(
            &data,
            count,
            dimension,
            DISKANN_MAX_PQ_TRAINING_VECTORS,
            74,
        )
        .unwrap();

        assert_eq!(first.len(), DISKANN_MAX_PQ_TRAINING_VECTORS * dimension);
        assert_eq!(first, second);
        assert_ne!(first, different_seed);
        assert_ne!(first, data[..DISKANN_MAX_PQ_TRAINING_VECTORS * dimension]);
        assert!(bounded_pq_training_sample(
            &data[..DISKANN_MAX_PQ_TRAINING_VECTORS * dimension],
            DISKANN_MAX_PQ_TRAINING_VECTORS,
            dimension,
            DISKANN_MAX_PQ_TRAINING_VECTORS,
            73
        )
        .is_none());
    }

    #[test]
    fn diskann_pq_training_plan_throttles_parallelism_and_samples_to_budget() {
        let unrestricted =
            pq_training_plan(1024, 256, 256, 50_000, 8 * 1024 * 1024 * 1024).unwrap();
        assert_eq!(unrestricted.sample_count, 50_000);
        assert_eq!(
            unrestricted.parallelism,
            rayon::current_num_threads().min(256)
        );

        let bounded = pq_training_plan(1024, 256, 256, 50_000, 64 * 1024 * 1024).unwrap();
        assert!(bounded.sample_count < unrestricted.sample_count);
        assert!(bounded.parallelism <= unrestricted.parallelism);
        assert!(pq_training_plan(1024, 1, 256, 256, 1).is_err());
    }

    #[test]
    fn diskann_pq_training_plan_accounts_for_retained_and_normalized_samples() {
        let budget = 128 * 1024 * 1024;
        let retained =
            pq_training_plan_with_sample_buffers(1024, 256, 256, 50_000, budget, 1).unwrap();
        let cosine =
            pq_training_plan_with_sample_buffers(1024, 256, 256, 50_000, budget, 2).unwrap();
        assert!(retained.sample_count < 50_000);
        assert!(cosine.sample_count < retained.sample_count);
    }

    #[test]
    fn low_level_diskann_training_uses_metric_aware_budget_and_propagates_failure() {
        let budget = 128 * 1024 * 1024;
        let build_params = DiskAnnBuildParams {
            memory_budget_bytes: budget,
            ..DiskAnnBuildParams::default()
        };
        let l2 = DiskAnnIndex::new(1024, MetricType::L2, 256, build_params);
        let cosine = DiskAnnIndex::new(1024, MetricType::Cosine, 256, build_params);
        let l2_plan = l2.training_plan(50_000).unwrap();
        let cosine_plan = cosine.training_plan(50_000).unwrap();

        assert_eq!(
            l2_plan,
            pq_training_plan_with_sample_buffers(1024, 256, 256, 50_000, budget, 1).unwrap()
        );
        assert_eq!(
            cosine_plan,
            pq_training_plan_with_sample_buffers(1024, 256, 256, 50_000, budget, 2).unwrap()
        );
        assert!(cosine_plan.sample_count < l2_plan.sample_count);

        let mut infeasible = DiskAnnIndex::new(
            8,
            MetricType::Cosine,
            2,
            DiskAnnBuildParams {
                memory_budget_bytes: 1,
                ..DiskAnnBuildParams::default()
            },
        );
        let error = infeasible
            .train(&[0.0; 8], 1)
            .expect_err("an infeasible low-level training budget must be returned");
        assert!(error.to_string().contains("PQ training peak"));
    }

    #[test]
    fn diskann_index_rejects_build_exceeding_memory_budget() {
        let build_params = DiskAnnBuildParams {
            memory_budget_bytes: 1,
            ..DiskAnnBuildParams::default()
        };
        let mut index = DiskAnnIndex::new(8, MetricType::L2, 2, build_params);
        index.add(&[0.0; 8], &[7]);

        let error = index
            .graph_build_shard_count()
            .expect_err("build should fail before graph allocation");
        assert!(error.to_string().contains("memory budget"));
    }

    #[test]
    fn diskann_memory_estimate_accounts_for_one_compact_graph_during_remap() {
        let count = 1024;
        let max_degree = 1024;
        let estimate = estimate_vamana_memory_bytes(count, max_degree, 1024, 1)
            .unwrap()
            .remap_peak_bytes;
        let one_graph = count * (max_degree * size_of::<u32>() + size_of::<u16>());

        assert!(
            estimate >= one_graph && estimate < 2 * one_graph,
            "in-place remapping should retain one compact graph; estimate={estimate}, one_graph={one_graph}"
        );
    }

    #[test]
    fn diskann_memory_estimate_reserves_row_id_encoding_scratch() {
        assert_eq!(row_id_encoding_scratch_bytes(1024).unwrap(), 8 * 1024);
        assert!(row_id_encoding_scratch_bytes(usize::MAX).is_err());
    }

    #[test]
    fn diskann_memory_budget_automatically_selects_overlapping_shards() {
        // Pin the worker count so this test exercises the same memory-plan
        // relationship on small CI runners and developer machines.
        rayon::ThreadPoolBuilder::new()
            .num_threads(12)
            .build()
            .unwrap()
            .install(|| {
                let count = 2_048;
                let dimension = 8;
                let params = DiskAnnBuildParams {
                    max_degree: 256,
                    build_search_list_size: 256,
                    ..DiskAnnBuildParams::default()
                };
                let mut index = DiskAnnIndex::new(dimension, MetricType::L2, 2, params);
                index.ids = (0..count as i64).collect();
                index.vectors = vec![0.0; count * dimension];
                let full_peak = index.estimate_build_memory_bytes().unwrap();
                index.build_params.memory_budget_bytes = full_peak - 1;

                let shard_count = index.graph_build_shard_count().unwrap();

                assert!(shard_count > 1);
                assert!(shard_count <= 64);
            });
    }

    #[test]
    fn diskann_locality_permutation_uses_bfs_then_unreachable_ids() {
        let graph =
            VamanaGraph::from_adjacency(2, vec![vec![], vec![], vec![3, 1], vec![4], vec![]]);

        let permutation = bfs_locality_permutation(&graph);

        assert_eq!(permutation.new_to_old, vec![2, 3, 1, 4, 0]);
        assert_eq!(permutation.old_to_new, vec![4, 2, 0, 1, 3]);
    }

    #[test]
    fn diskann_locality_remaps_entry_and_neighbor_ids() {
        let mut graph =
            VamanaGraph::from_adjacency(2, vec![vec![], vec![], vec![3, 1], vec![4], vec![]]);
        let permutation = bfs_locality_permutation(&graph);

        remap_graph_in_place(&mut graph, &permutation);

        assert_eq!(graph.entry_node, 0);
        assert_eq!(
            graph
                .adjacency
                .iter()
                .map(<[u32]>::to_vec)
                .collect::<Vec<_>>(),
            vec![vec![1, 2], vec![3], vec![], vec![], vec![]]
        );
    }

    #[test]
    fn diskann_prepare_build_keeps_rows_codes_and_vectors_aligned() {
        let dimension = 8;
        let training_count = 256;
        let indexed_count = 64;
        let data = (0..training_count * dimension)
            .map(|offset| ((offset * 29) % 113) as f32)
            .collect::<Vec<_>>();
        let ids = (1000..1000 + indexed_count as i64).collect::<Vec<_>>();
        let mut index = DiskAnnIndex::new(
            dimension,
            MetricType::L2,
            2,
            DiskAnnBuildParams {
                max_degree: 8,
                build_search_list_size: 16,
                ..DiskAnnBuildParams::default()
            },
        );
        index.train(&data, training_count).unwrap();
        index.add(&data[..indexed_count * dimension], &ids);

        let prepared = index.prepare_build().unwrap();

        assert_eq!(prepared.graph.entry_node, 0);
        assert_eq!(prepared.pq_codes.len(), indexed_count * 2);
        for new_id in 0..indexed_count {
            let old_id = prepared.permutation.new_to_old[new_id] as usize;
            assert_eq!(prepared.row_id(&index, new_id), ids[old_id]);
            assert_eq!(
                prepared.vector(&index, new_id),
                &data[old_id * dimension..(old_id + 1) * dimension]
            );
            assert_eq!(prepared.pq_code(&index, new_id).len(), 2);
        }
    }

    #[test]
    fn diskann_cosine_normalizes_vectors_before_pq_encoding_and_persistence() {
        let mut index =
            DiskAnnIndex::with_pq_bits(2, MetricType::Cosine, 1, 4, DiskAnnBuildParams::default());

        index.add(&[3.0, 4.0, 0.0, 0.0], &[10, 11]);

        assert_eq!(index.vectors, vec![0.6, 0.8, 0.0, 0.0]);
    }

    #[test]
    fn diskann_only_allocates_metric_preprocessing_for_cosine() {
        let data = [3.0, 4.0];
        for metric in [MetricType::L2, MetricType::InnerProduct] {
            let index = DiskAnnIndex::with_pq_bits(2, metric, 1, 4, DiskAnnBuildParams::default());
            assert!(matches!(
                index.preprocess_vectors(&data, 1),
                Cow::Borrowed(_)
            ));
        }
        let index =
            DiskAnnIndex::with_pq_bits(2, MetricType::Cosine, 1, 4, DiskAnnBuildParams::default());
        assert!(matches!(index.preprocess_vectors(&data, 1), Cow::Owned(_)));
    }
}
