use std::borrow::Borrow;
use std::path::{Path, PathBuf};

use common::bitvec::{BitSlice, DeletedBitVec};
use common::fs::{atomic_save_json, clear_disk_cache};
use common::mmap::{AdviceSetting, MmapSlice, create_and_ensure_length};
use common::types::PointOffsetType;
use common::universal_io::{
    CachedReadFs, OkNotFound, OpenOptions, Populate, SortedBlockIndex, TypedStorage, UniversalRead,
    UniversalReadFs, read_json_via,
};
use fs_err as fs;
use memmap2::MmapMut;
use serde::{Deserialize, Serialize};

use super::super::Encodable;
use super::super::mutable_numeric_index::InMemoryNumericIndex;
use super::{
    CONFIG_PATH, DELETED_PATH, OnDiskNumericIndex, PAIRS_BLOCK_INDEX_PATH, PAIRS_PATH, Storage,
};
use crate::common::Flusher;
use crate::common::operation_error::{OperationError, OperationResult};
use crate::index::field_index::deleted_mask::{
    bitor_deleted_mask, deleted_mask_file, preopen_deleted_mask, save_deleted_mask,
};
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::numeric_point::{Numericable, Point};
use crate::index::field_index::on_disk_point_to_values::{OnDiskPointToValues, StoredValue};

#[derive(Debug, Clone, Serialize, Deserialize)]
struct UniversalNumericIndexConfig {
    max_values_per_point: usize,
}

impl<T, S> OnDiskNumericIndex<T, S>
where
    T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod,
    S: UniversalRead,
{
    /// TODO: save using `S::Fs` too
    pub fn build(
        fs: &S::Fs,
        in_memory_index: InMemoryNumericIndex<T>,
        path: &Path,
        populate: Populate,
        deleted_points: &BitSlice,
    ) -> OperationResult<Self> {
        fs::create_dir_all(path)?;

        let pairs_path = path.join(PAIRS_PATH);
        let config_path = path.join(CONFIG_PATH);

        atomic_save_json(
            &config_path,
            &UniversalNumericIndexConfig {
                max_values_per_point: in_memory_index.max_values_per_point,
            },
        )?;

        in_memory_index.histogram.save(path)?;

        OnDiskPointToValues::<T, S>::build_from_iter(
            path,
            in_memory_index
                .point_to_values
                .iter()
                .enumerate()
                .map(|(idx, values)| (idx as PointOffsetType, values.iter().map(|v| v.borrow()))),
        )?;

        {
            let pairs_file = create_and_ensure_length(
                &pairs_path,
                in_memory_index.map.len() * size_of::<Point<T>>(),
            )?;
            let pairs_mmap = unsafe { MmapMut::map_mut(&pairs_file)? };
            let mut pairs = unsafe { MmapSlice::<Point<T>>::try_from(pairs_mmap)? };
            for (src, dst) in in_memory_index.map.iter().zip(pairs.iter_mut()) {
                *dst = *src;
            }

            SortedBlockIndex::write(&path.join(PAIRS_BLOCK_INDEX_PATH), &pairs)?;
        }

        save_deleted_mask(
            path,
            DELETED_PATH,
            in_memory_index.point_to_values.len(),
            in_memory_index
                .point_to_values
                .iter()
                .enumerate()
                .filter(|(_, values)| values.is_empty())
                .map(|(idx, _)| idx as PointOffsetType),
        )?;

        Self::open(fs, path, populate, deleted_points)?.ok_or_else(|| {
            OperationError::service_error("Failed to open UniversalNumericIndex after building it")
        })
    }

    fn open_options(populate: Populate) -> OpenOptions {
        OpenOptions {
            writeable: false,
            need_sequential: false,
            populate,
            advice: AdviceSetting::Global,
        }
    }

    /// Schedule background prefetch of every file [`open`](Self::open) will read.
    ///
    /// Returns `false` (nothing scheduled) when the segment is not in the
    /// on-disk format.
    pub fn preopen(
        fs: &impl CachedReadFs<File = S>,
        path: &Path,
        populate: Populate,
    ) -> OperationResult<bool> {
        // Config
        let config_path = path.join(CONFIG_PATH);
        if fs
            .schedule_prefetch(&config_path, None, None)
            .ok_not_found()?
            .is_none()
        {
            // If config doesn't exist, assume the index doesn't exist on disk
            return Ok(false);
        }

        // Histogram
        Histogram::<T>::preopen(fs, path)?;

        // Value pairs
        let pairs_path = path.join(PAIRS_PATH);
        fs.schedule_prefetch(&pairs_path, Some(Self::open_options(populate)), None)?;

        // Block index over the value pairs; optional, absent on old segments
        let _ = fs
            .schedule_prefetch(&path.join(PAIRS_BLOCK_INDEX_PATH), None, None)
            .ok_not_found()?;

        // Point to values
        OnDiskPointToValues::<T, S>::preopen(fs, path, populate)?;

        // "No values" mask
        preopen_deleted_mask(
            fs,
            path,
            DELETED_PATH,
            Self::open_options(Populate::PreferBackground),
        )?;

        Ok(true)
    }

    /// Open and load mmap numeric index from the given path
    pub fn open(
        fs: &impl UniversalReadFs<File = S>,
        path: &Path,
        populate: Populate,
        deleted_points: &BitSlice,
    ) -> OperationResult<Option<Self>> {
        let pairs_path = path.join(PAIRS_PATH);
        let config_path = path.join(CONFIG_PATH);

        let Some(config) =
            read_json_via::<_, UniversalNumericIndexConfig>(fs, &config_path).ok_not_found()?
        else {
            // If config doesn't exist, assume the index doesn't exist on disk
            return Ok(None);
        };

        let histogram = Histogram::<T>::open(fs, path)?;

        let pairs = TypedStorage::new(fs.open(
            &pairs_path,
            Self::open_options(populate),
            Default::default(),
        )?);

        let pairs_block_index = SortedBlockIndex::open(
            fs,
            &path.join(PAIRS_BLOCK_INDEX_PATH),
            pairs.len()? as usize,
        )?;

        let point_to_values = OnDiskPointToValues::open(fs, path, populate)?;
        let mut deleted = deleted_points.to_owned();

        // `deleted` length must match `point_to_values.len()` because it only
        // tracks the index's contents. The id-tracker's deleted mask can be
        // shorter or longer; if shorter, the missing entries default to live
        // (the id-tracker is the source of truth for deletions, and a shorter
        // mask just means it doesn't yet know about those higher offsets).
        deleted.resize(point_to_values.len(), false);
        let compact_deleted_mask = bitor_deleted_mask(
            fs,
            path,
            DELETED_PATH,
            Self::open_options(Populate::Auto),
            &mut deleted,
        )?;

        Ok(Some(Self {
            path: path.to_path_buf(),
            storage: Storage {
                deleted: DeletedBitVec::new(deleted),
                pairs,
                pairs_block_index,
                point_to_values,
            },
            histogram,
            max_values_per_point: config.max_values_per_point,
            compact_deleted_mask,
        }))
    }
}

impl<T, S> OnDiskNumericIndex<T, S>
where
    T: Encodable + Numericable + Default + StoredValue + 'static,
    S: UniversalRead,
{
    pub fn wipe(self) -> OperationResult<()> {
        let files = self.files();
        let path = self.path.clone();
        // drop mmap handles before deleting files
        drop(self);
        for file in files {
            fs::remove_file(file)?;
        }
        let _ = fs::remove_dir(path);
        Ok(())
    }

    pub fn files(&self) -> Vec<PathBuf> {
        let mut files = vec![
            self.path.join(PAIRS_PATH),
            deleted_mask_file(&self.path, self.compact_deleted_mask, DELETED_PATH),
            self.path.join(CONFIG_PATH),
        ];
        if self.storage.pairs_block_index.is_some() {
            files.push(self.path.join(PAIRS_BLOCK_INDEX_PATH));
        }
        files.extend(self.storage.point_to_values.files());
        files.extend(Histogram::<T>::files(&self.path));
        files
    }

    pub fn immutable_files(&self) -> Vec<PathBuf> {
        let mut files = vec![
            self.path.join(PAIRS_PATH),
            deleted_mask_file(&self.path, self.compact_deleted_mask, DELETED_PATH),
            self.path.join(CONFIG_PATH),
        ];
        if self.storage.pairs_block_index.is_some() {
            files.push(self.path.join(PAIRS_BLOCK_INDEX_PATH));
        }
        files.extend(self.storage.point_to_values.immutable_files());
        files.extend(Histogram::<T>::immutable_files(&self.path));
        files
    }

    /// No-op flusher: the on-disk state is build-time only. See the type-level
    /// docs on [`UniversalNumericIndex`] for the deletion durability contract.
    pub fn flusher(&self) -> Flusher {
        Box::new(|| Ok(()))
    }

    /// Marks `idx` as deleted in the in-memory deletion bitvec.
    ///
    /// Not persisted: on reopen, deletions must be re-supplied via the
    /// `deleted_points` argument to [`Self::open`].
    pub fn remove_point(&mut self, idx: PointOffsetType) {
        self.storage.deleted.mark_deleted(idx);
    }

    /// Populate all pages in the mmap.
    /// Block until all pages are populated.
    pub fn populate(&self) -> OperationResult<()> {
        self.storage.pairs.populate()?;
        self.storage.point_to_values.populate()?;
        Ok(())
    }

    /// Drop disk cache.
    pub fn clear_cache(&self) -> OperationResult<()> {
        let Self {
            path,
            storage,
            histogram: _,
            max_values_per_point: _,
            compact_deleted_mask,
        } = self;
        let Storage {
            deleted: _,
            pairs,
            pairs_block_index,
            point_to_values,
        } = storage;
        pairs.clear_ram_cache()?;
        if pairs_block_index.is_some() {
            clear_disk_cache(&path.join(PAIRS_BLOCK_INDEX_PATH))?;
        }
        clear_disk_cache(&deleted_mask_file(
            path,
            *compact_deleted_mask,
            DELETED_PATH,
        ))?;
        point_to_values.clear_cache()?;
        Ok(())
    }

    pub(crate) fn ram_usage_bytes(&self) -> usize {
        let Self {
            path: _,
            storage,
            histogram,
            max_values_per_point: _,
            compact_deleted_mask: _,
        } = self;

        histogram.ram_usage_bytes() + storage.ram_usage_bytes()
    }
}
