use std::borrow::Cow;
use std::fmt;
use std::ops::Range;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use bitvec::prelude::{BitSlice, BitVec};
use common::counter::hardware_counter::HardwareCounterCell;
use common::ext::BitSliceExt as _;
use common::types::PointOffsetType;
use parking_lot::RwLock;
use rocksdb::DB;

use crate::common::Flusher;
use crate::common::operation_error::{OperationError, OperationResult, check_process_stopped};
use crate::common::rocksdb_wrapper::DatabaseColumnWrapper;
use crate::data_types::named_vectors::{CowMultiVector, CowVector};
use crate::data_types::primitive::PrimitiveVectorElement;
use crate::data_types::vectors::{
    TypedMultiDenseVector, TypedMultiDenseVectorRef, VectorElementType, VectorRef,
};
use crate::types::{Distance, MultiVectorConfig, VectorStorageDatatype};
use crate::vector_storage::bitvec::bitvec_set_deleted;
use crate::vector_storage::common::{CHUNK_SIZE, StoredRecord};
use crate::vector_storage::volatile_chunked_vectors::VolatileChunkedVectors;
use crate::vector_storage::{
    AccessPattern, MultiVectorStorage, VectorOffsetType, VectorStorage, VectorStorageEnum,
};

type StoredMultiDenseVector<T> = StoredRecord<TypedMultiDenseVector<T>>;

/// All fields are counting vectors and not dimensions.
#[derive(Debug, Clone, Default)]
struct MultiVectorMetadata {
    id: VectorOffsetType,
    start: VectorOffsetType,
    inner_vectors_count: usize,
    inner_vector_capacity: usize,
}

/// In-memory vector storage with on-update persistence using `store`
pub struct SimpleMultiDenseVectorStorage<T: PrimitiveVectorElement> {
    dim: usize,
    distance: Distance,
    multi_vector_config: MultiVectorConfig,
    /// Keep vectors in memory
    vectors: VolatileChunkedVectors<T>,
    vectors_metadata: Vec<MultiVectorMetadata>,
    db_wrapper: DatabaseColumnWrapper,
    /// BitVec for deleted flags. Grows dynamically upto last set flag.
    deleted: BitVec,
    /// Current number of deleted vectors.
    deleted_count: usize,
}

impl<T: fmt::Debug + PrimitiveVectorElement> fmt::Debug for SimpleMultiDenseVectorStorage<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SimpleMultiDenseVectorStorage")
            .field("dim", &self.dim)
            .field("distance", &self.distance)
            .field("multi_vector_config", &self.multi_vector_config)
            .field("vectors", &self.vectors)
            .field("vectors_metadata", &self.vectors_metadata)
            .field("db_wrapper", &self.db_wrapper)
            .field("deleted_count", &self.deleted_count)
            .finish_non_exhaustive()
    }
}

pub fn open_simple_multi_dense_vector_storage(
    storage_element_type: VectorStorageDatatype,
    database: Arc<RwLock<DB>>,
    database_column_name: &str,
    dim: usize,
    distance: Distance,
    multi_vector_config: MultiVectorConfig,
    stopped: &AtomicBool,
) -> OperationResult<VectorStorageEnum> {
    match storage_element_type {
        VectorStorageDatatype::Float32 => open_simple_multi_dense_vector_storage_full(
            database,
            database_column_name,
            dim,
            distance,
            multi_vector_config,
            stopped,
        ),
        VectorStorageDatatype::Uint8 => open_simple_multi_dense_vector_storage_byte(
            database,
            database_column_name,
            dim,
            distance,
            multi_vector_config,
            stopped,
        ),
        VectorStorageDatatype::Float16 => open_simple_multi_dense_vector_storage_half(
            database,
            database_column_name,
            dim,
            distance,
            multi_vector_config,
            stopped,
        ),
    }
}

pub fn open_simple_multi_dense_vector_storage_full(
    database: Arc<RwLock<DB>>,
    database_column_name: &str,
    dim: usize,
    distance: Distance,
    multi_vector_config: MultiVectorConfig,
    stopped: &AtomicBool,
) -> OperationResult<VectorStorageEnum> {
    let storage = open_simple_multi_dense_vector_storage_impl(
        database,
        database_column_name,
        dim,
        distance,
        multi_vector_config,
        stopped,
    )?;
    Ok(VectorStorageEnum::MultiDenseSimple(storage))
}

pub fn open_simple_multi_dense_vector_storage_byte(
    database: Arc<RwLock<DB>>,
    database_column_name: &str,
    dim: usize,
    distance: Distance,
    multi_vector_config: MultiVectorConfig,
    stopped: &AtomicBool,
) -> OperationResult<VectorStorageEnum> {
    let storage = open_simple_multi_dense_vector_storage_impl(
        database,
        database_column_name,
        dim,
        distance,
        multi_vector_config,
        stopped,
    )?;
    Ok(VectorStorageEnum::MultiDenseSimpleByte(storage))
}

pub fn open_simple_multi_dense_vector_storage_half(
    database: Arc<RwLock<DB>>,
    database_column_name: &str,
    dim: usize,
    distance: Distance,
    multi_vector_config: MultiVectorConfig,
    stopped: &AtomicBool,
) -> OperationResult<VectorStorageEnum> {
    let storage = open_simple_multi_dense_vector_storage_impl(
        database,
        database_column_name,
        dim,
        distance,
        multi_vector_config,
        stopped,
    )?;
    Ok(VectorStorageEnum::MultiDenseSimpleHalf(storage))
}

fn open_simple_multi_dense_vector_storage_impl<T: PrimitiveVectorElement>(
    database: Arc<RwLock<DB>>,
    database_column_name: &str,
    dim: usize,
    distance: Distance,
    multi_vector_config: MultiVectorConfig,
    stopped: &AtomicBool,
) -> OperationResult<SimpleMultiDenseVectorStorage<T>> {
    let mut vectors = VolatileChunkedVectors::new(dim);
    let mut vectors_metadata = Vec::<MultiVectorMetadata>::new();
    let (mut deleted, mut deleted_count) = (BitVec::new(), 0);
    let db_wrapper = DatabaseColumnWrapper::new(database, database_column_name);
    db_wrapper.lock_db().iter()?;
    for (key, value) in db_wrapper.lock_db().iter()? {
        let point_id: PointOffsetType = bincode::deserialize(&key)
            .map_err(|_| OperationError::service_error("cannot deserialize point id from db"))?;
        let stored_record: StoredMultiDenseVector<T> = bincode::deserialize(&value)
            .map_err(|_| OperationError::service_error("cannot deserialize record from db"))?;

        // Propagate deleted flag
        if stored_record.deleted {
            bitvec_set_deleted(&mut deleted, point_id, true);
            deleted_count += 1;
        }
        let point_id_usize = point_id as usize;
        if point_id_usize >= vectors_metadata.len() {
            vectors_metadata.resize(point_id_usize + 1, Default::default());
        }

        let metadata = &mut vectors_metadata[point_id_usize];
        metadata.inner_vectors_count = stored_record.vector.vectors_count();
        metadata.inner_vector_capacity = metadata.inner_vectors_count;
        metadata.id = point_id as VectorOffsetType;

        metadata.start = vectors.len();
        let left_keys = vectors.get_chunk_left_keys(metadata.start);
        if stored_record.vector.vectors_count() > left_keys {
            metadata.start += left_keys;
        }
        vectors.insert_many(
            metadata.start,
            &stored_record.vector.flattened_vectors,
            stored_record.vector.vectors_count(),
        )?;

        check_process_stopped(stopped)?;
    }

    Ok(SimpleMultiDenseVectorStorage {
        dim,
        distance,
        multi_vector_config,
        vectors,
        vectors_metadata,
        db_wrapper,
        deleted,
        deleted_count,
    })
}

impl<T: PrimitiveVectorElement> SimpleMultiDenseVectorStorage<T> {
    /// Set deleted flag for given key. Returns previous deleted state.
    #[inline]
    fn set_deleted(&mut self, key: PointOffsetType, deleted: bool) -> bool {
        if !deleted && key as usize >= self.vectors.len() {
            return false;
        }
        let was_deleted = bitvec_set_deleted(&mut self.deleted, key, deleted);
        if was_deleted != deleted {
            if !was_deleted {
                self.deleted_count += 1;
            } else {
                self.deleted_count = self.deleted_count.saturating_sub(1);
            }
        }
        was_deleted
    }

    fn update_stored(
        &self,
        key: PointOffsetType,
        deleted: bool,
        vector: Option<TypedMultiDenseVectorRef<'_, T>>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        let mut record = StoredMultiDenseVector {
            deleted,
            vector: TypedMultiDenseVector::placeholder(self.dim),
        };
        if let Some(vector) = vector {
            record.vector.dim = vector.dim;
            record.vector.flattened_vectors.clear();
            record
                .vector
                .flattened_vectors
                .extend_from_slice(vector.flattened_vectors);
        }

        let key_enc = bincode::serialize(&key).unwrap();
        let record_enc = bincode::serialize(&record).unwrap();

        hw_counter
            .vector_io_write_counter()
            .incr_delta(key_enc.len() + record_enc.len());

        // Store updated record
        self.db_wrapper.put(key_enc, record_enc)?;

        Ok(())
    }

    fn insert_vector_impl(
        &mut self,
        key: PointOffsetType,
        vector: VectorRef,
        is_deleted: bool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        let multi_vector: TypedMultiDenseVectorRef<VectorElementType> = vector.try_into()?;
        let multi_vector = T::from_float_multivector(CowMultiVector::Borrowed(multi_vector));
        let multi_vector = multi_vector.as_vec_ref();
        assert_eq!(multi_vector.dim, self.dim);
        let multivector_size_in_bytes = std::mem::size_of_val(multi_vector.flattened_vectors);
        if multivector_size_in_bytes >= CHUNK_SIZE {
            return Err(OperationError::service_error(format!(
                "Cannot insert multi vector of size {multivector_size_in_bytes} to the vector storage. It's too large, maximum size is {CHUNK_SIZE}.",
            )));
        }

        let key_usize = key as usize;
        if key_usize >= self.vectors_metadata.len() {
            self.vectors_metadata
                .resize(key_usize + 1, Default::default());
        }
        let metadata = &mut self.vectors_metadata[key_usize];
        metadata.id = key as VectorOffsetType;
        metadata.inner_vectors_count = multi_vector.vectors_count();

        if multi_vector.vectors_count() > metadata.inner_vector_capacity {
            metadata.inner_vector_capacity = metadata.inner_vectors_count;
            metadata.start = self.vectors.len();
            let left_keys = self.vectors.get_chunk_left_keys(metadata.start);
            if multi_vector.vectors_count() > left_keys {
                metadata.start += left_keys;
            }
            self.vectors.insert_many(
                metadata.start,
                multi_vector.flattened_vectors,
                multi_vector.vectors_count(),
            )?;
        } else {
            self.vectors.insert_many(
                metadata.start,
                multi_vector.flattened_vectors,
                multi_vector.vectors_count(),
            )?;
        }

        self.set_deleted(key, is_deleted);
        self.update_stored(key, is_deleted, Some(multi_vector), hw_counter)?;
        Ok(())
    }

    /// Destroy this vector storage, remove persisted data from RocksDB
    pub fn destroy(&self) -> OperationResult<()> {
        self.db_wrapper.remove_column_family()?;
        Ok(())
    }
}

impl<T: PrimitiveVectorElement> MultiVectorStorage<T> for SimpleMultiDenseVectorStorage<T> {
    fn vector_dim(&self) -> usize {
        self.dim
    }

    /// Panics if key is out of bounds
    fn get_multi<P: AccessPattern>(&self, key: PointOffsetType) -> CowMultiVector<'_, T> {
        self.get_multi_opt::<P>(key).expect("vector not found")
    }

    /// None if key is out of bounds
    fn get_multi_opt<P: AccessPattern>(
        &self,
        key: PointOffsetType,
    ) -> Option<CowMultiVector<'_, T>> {
        // No sequential optimizations available for in memory storage.
        self.vectors_metadata.get(key as usize).map(|metadata| {
            let flattened_vectors = self
                .vectors
                .get_many(metadata.start, metadata.inner_vectors_count)
                .unwrap_or_else(|| panic!("Vectors does not contain data for {metadata:?}"));
            CowMultiVector::Borrowed(TypedMultiDenseVectorRef {
                flattened_vectors,
                dim: self.dim,
            })
        })
    }

    fn iterate_inner_vectors(&self) -> impl Iterator<Item = Cow<'_, [T]>> + Clone + Send {
        (0..self.total_vector_count()).flat_map(|key| {
            let metadata = &self.vectors_metadata[key];
            (0..metadata.inner_vectors_count)
                .map(|i| Cow::Borrowed(self.vectors.get(metadata.start + i)))
        })
    }

    fn multi_vector_config(&self) -> &MultiVectorConfig {
        &self.multi_vector_config
    }

    fn size_of_available_vectors_in_bytes(&self) -> usize {
        if self.total_vector_count() > 0 {
            let total_size = self.vectors.len() * self.vector_dim() * std::mem::size_of::<T>();
            (total_size as u128 * self.available_vector_count() as u128
                / self.total_vector_count() as u128) as usize
        } else {
            0
        }
    }
}

impl<T: PrimitiveVectorElement> VectorStorage for SimpleMultiDenseVectorStorage<T> {
    fn distance(&self) -> Distance {
        self.distance
    }

    fn datatype(&self) -> VectorStorageDatatype {
        VectorStorageDatatype::Float32
    }

    fn is_on_disk(&self) -> bool {
        false
    }

    fn total_vector_count(&self) -> usize {
        self.vectors_metadata.len()
    }

    fn get_vector<P: AccessPattern>(&self, key: PointOffsetType) -> CowVector<'_> {
        self.get_vector_opt::<P>(key).expect("vector not found")
    }

    fn get_vector_opt<P: AccessPattern>(&self, key: PointOffsetType) -> Option<CowVector<'_>> {
        self.get_multi_opt::<P>(key).map(|multi_dense_vector| {
            CowVector::MultiDense(T::into_float_multivector(multi_dense_vector))
        })
    }

    fn insert_vector(
        &mut self,
        key: PointOffsetType,
        vector: VectorRef,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        self.insert_vector_impl(key, vector, false, hw_counter)
    }

    fn update_from<'a>(
        &mut self,
        other_vectors: &'a mut impl Iterator<Item = (CowVector<'a>, bool)>,
        stopped: &AtomicBool,
    ) -> OperationResult<Range<PointOffsetType>> {
        let start_index = self.vectors_metadata.len() as PointOffsetType;
        for (other_vector, other_deleted) in other_vectors {
            check_process_stopped(stopped)?;
            // Do not perform preprocessing - vectors should be already processed
            let other_vector: VectorRef = other_vector.as_vec_ref();
            let new_id = self.vectors_metadata.len() as PointOffsetType;
            self.insert_vector_impl(
                new_id,
                other_vector,
                other_deleted,
                &HardwareCounterCell::disposable(), // This function is only used by internal operations
            )?;
        }
        let end_index = self.vectors_metadata.len() as PointOffsetType;
        Ok(start_index..end_index)
    }

    fn flusher(&self) -> Flusher {
        self.db_wrapper.flusher()
    }

    fn files(&self) -> Vec<std::path::PathBuf> {
        vec![]
    }

    fn delete_vector(&mut self, key: PointOffsetType) -> OperationResult<bool> {
        let is_deleted = !self.set_deleted(key, true);
        if is_deleted {
            // We don't measure deletions.
            self.update_stored(key, true, None, &HardwareCounterCell::disposable())?;
        }
        Ok(is_deleted)
    }

    fn is_deleted_vector(&self, key: PointOffsetType) -> bool {
        self.deleted.get_bit(key as usize).unwrap_or(false)
    }

    fn deleted_vector_count(&self) -> usize {
        self.deleted_count
    }

    fn deleted_vector_bitslice(&self) -> &BitSlice {
        self.deleted.as_bitslice()
    }
}

#[cfg(test)]
mod tests {
    use rand::rngs::StdRng;
    use rand::{RngExt, SeedableRng};
    use tempfile::Builder;

    use super::*;
    use crate::common::rocksdb_wrapper::{DB_VECTOR_CF, open_db};
    use crate::data_types::vectors::MultiDenseVectorInternal;
    use crate::segment_constructor::migrate_rocksdb_multi_dense_vector_storage_to_mmap;
    use crate::vector_storage::Sequential;

    const RAND_SEED: u64 = 42;

    /// Create RocksDB based multi dense vector storage.
    ///
    /// Migrate it to the mmap based multi dense vector storage and assert vector data is correct.
    #[test]
    fn test_migrate_simple_to_mmap() {
        const POINT_COUNT: PointOffsetType = 128;
        const DIM: usize = 128;
        const DELETE_PROBABILITY: f64 = 0.1;

        let mut rng = StdRng::seed_from_u64(RAND_SEED);

        let multi_vector_config = MultiVectorConfig::default();
        let db_dir = Builder::new().prefix("storage_dir").tempdir().unwrap();
        let db = open_db(db_dir.path(), &[DB_VECTOR_CF]).unwrap();

        // Create simple multi dense vector storage, insert test points and delete some of them again
        let mut storage = open_simple_multi_dense_vector_storage_full(
            db,
            DB_VECTOR_CF,
            DIM,
            Distance::Dot,
            multi_vector_config,
            &AtomicBool::new(false),
        )
        .unwrap();
        for internal_id in 0..POINT_COUNT {
            let size = rng.random_range(1..=4);
            let vectors = std::iter::repeat_with(|| {
                std::iter::repeat_with(|| rng.random_range(-1.0..1.0))
                    .take(DIM)
                    .collect()
            })
            .take(size)
            .collect::<Vec<Vec<_>>>();
            let multivec = MultiDenseVectorInternal::try_from(vectors).unwrap();
            storage
                .insert_vector(
                    internal_id,
                    VectorRef::from(&multivec),
                    &HardwareCounterCell::disposable(),
                )
                .unwrap();
            if rng.random_bool(DELETE_PROBABILITY) {
                storage.delete_vector(internal_id).unwrap();
            }
        }

        let deleted_vector_count = storage.deleted_vector_count();
        let total_vector_count = storage.total_vector_count();

        // Migrate from RocksDB to mmap storage
        let storage_dir = Builder::new().prefix("storage_dir").tempdir().unwrap();
        let new_storage = migrate_rocksdb_multi_dense_vector_storage_to_mmap(
            &storage,
            DIM,
            multi_vector_config,
            storage_dir.path(),
        )
        .expect("failed to migrate from RocksDB to mmap");

        // Destroy persisted RocksDB dense vector data
        match storage {
            VectorStorageEnum::MultiDenseSimple(storage) => storage.destroy().unwrap(),
            VectorStorageEnum::MultiDenseSimpleByte(storage) => storage.destroy().unwrap(),
            VectorStorageEnum::MultiDenseSimpleHalf(storage) => storage.destroy().unwrap(),
            _ => unreachable!("unexpected vector storage type"),
        }

        // We can drop RocksDB storage now
        db_dir.close().expect("failed to drop RocksDB storage");

        // Assert vector counts and data
        let mut rng = StdRng::seed_from_u64(RAND_SEED);
        assert_eq!(new_storage.deleted_vector_count(), deleted_vector_count);
        assert_eq!(new_storage.total_vector_count(), total_vector_count);
        for internal_id in 0..POINT_COUNT {
            let size = rng.random_range(1..=4);
            let vectors = std::iter::repeat_with(|| {
                std::iter::repeat_with(|| rng.random_range(-1.0..1.0))
                    .take(DIM)
                    .collect()
            })
            .take(size)
            .collect::<Vec<Vec<_>>>();
            let multivec = MultiDenseVectorInternal::try_from(vectors).unwrap();
            assert_eq!(
                new_storage.get_vector::<Sequential>(internal_id),
                CowVector::from(&multivec),
            );
            assert_eq!(
                new_storage.is_deleted_vector(internal_id),
                rng.random_bool(DELETE_PROBABILITY)
            );
        }
    }
}
