use crate::blobstore::{
    BlobStore, BlobStoreCleanupStat, BlobStoreError, BlobStoreSize, PooledBlobSidecar,
};
use alloy_eips::{
    eip4844::{BlobAndProofV1, BlobAndProofV2, BlobCellsAndProofsV1},
    eip7594::{BlobCellMask, BlobTransactionSidecarVariant, Cell},
};
use alloy_primitives::{map::B256Map, B128, B256};
use parking_lot::RwLock;
use std::sync::Arc;

/// An in-memory blob store.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct InMemoryBlobStore {
    inner: Arc<InMemoryBlobStoreInner>,
}

impl InMemoryBlobStore {
    /// Look up EIP-7594 blobs by their versioned hashes.
    ///
    /// This returns a result vector with the **same length and order** as the input
    /// `versioned_hashes`. Each element is `Some(BlobAndProofV2)` if the blob is available, or
    /// `None` if it is missing or an older sidecar version.
    fn get_by_versioned_hashes_eip7594(
        &self,
        versioned_hashes: &[B256],
    ) -> Vec<Option<BlobAndProofV2>> {
        let mut result = vec![None; versioned_hashes.len()];
        let mut missing_count = result.len();
        for blob_sidecar in self.inner.store.read().values() {
            if let Some(blob_sidecar) = blob_sidecar.as_eip7594() {
                for (hash_idx, match_result) in
                    blob_sidecar.match_versioned_hashes(versioned_hashes)
                {
                    let slot = &mut result[hash_idx];
                    if slot.is_none() {
                        missing_count -= 1;
                    }
                    *slot = Some(match_result);
                }
            }

            // Return early if all blobs are found.
            if missing_count == 0 {
                // since versioned_hashes may have duplicates, we double check here
                if result.iter().all(|blob| blob.is_some()) {
                    break;
                }
            }
        }
        result
    }

    /// Look up EIP-7594 blob cells by their versioned hashes.
    fn get_by_versioned_hashes_cells_eip7594(
        &self,
        versioned_hashes: &[B256],
        indices_bitarray: B128,
    ) -> Result<Vec<Option<BlobCellsAndProofsV1>>, BlobStoreError> {
        let cell_mask = BlobCellMask::new(indices_bitarray);
        let mut result = vec![None; versioned_hashes.len()];
        let mut missing_count = result.len();
        let blob_sidecars = self.inner.store.read().values().cloned().collect::<Vec<_>>();
        for blob_sidecar in blob_sidecars {
            if let Some(blob_sidecar) = blob_sidecar.as_eip7594() {
                for (hash_idx, match_result) in blob_sidecar
                    .match_versioned_hashes_cells(versioned_hashes, cell_mask)
                    .map_err(|err| BlobStoreError::Other(Box::new(err)))?
                {
                    let slot = &mut result[hash_idx];
                    if slot.is_none() {
                        missing_count -= 1;
                    }
                    *slot = Some(match_result);
                }
            }

            if missing_count == 0 && result.iter().all(Option::is_some) {
                break;
            }
        }
        Ok(result)
    }
}

#[derive(Debug, Default)]
struct InMemoryBlobStoreInner {
    /// Storage for all blob data.
    store: RwLock<B256Map<Arc<BlobTransactionSidecarVariant>>>,
    size_tracker: BlobStoreSize,
}

impl PartialEq for InMemoryBlobStoreInner {
    fn eq(&self, other: &Self) -> bool {
        self.store.read().eq(&*other.store.read())
    }
}

impl BlobStore for InMemoryBlobStore {
    fn insert(&self, tx: B256, data: PooledBlobSidecar) -> Result<(), BlobStoreError> {
        let mut store = self.inner.store.write();
        self.inner.size_tracker.add_size(insert_size(&mut store, tx, data.into_sidecar()));
        self.inner.size_tracker.update_len(store.len());
        Ok(())
    }

    fn insert_all(&self, txs: Vec<(B256, PooledBlobSidecar)>) -> Result<(), BlobStoreError> {
        if txs.is_empty() {
            return Ok(())
        }
        let mut store = self.inner.store.write();
        let mut total_add = 0;
        for (tx, data) in txs {
            let add = insert_size(&mut store, tx, data.into_sidecar());
            total_add += add;
        }
        self.inner.size_tracker.add_size(total_add);
        self.inner.size_tracker.update_len(store.len());
        Ok(())
    }

    fn delete(&self, tx: B256) -> Result<(), BlobStoreError> {
        let mut store = self.inner.store.write();
        let sub = remove_size(&mut store, &tx);
        self.inner.size_tracker.sub_size(sub);
        self.inner.size_tracker.update_len(store.len());
        Ok(())
    }

    fn delete_all(&self, txs: Vec<B256>) -> Result<(), BlobStoreError> {
        if txs.is_empty() {
            return Ok(())
        }
        let mut store = self.inner.store.write();
        let mut total_sub = 0;
        for tx in txs {
            total_sub += remove_size(&mut store, &tx);
        }
        self.inner.size_tracker.sub_size(total_sub);
        self.inner.size_tracker.update_len(store.len());
        Ok(())
    }

    fn cleanup(&self) -> BlobStoreCleanupStat {
        BlobStoreCleanupStat::default()
    }

    // Retrieves the decoded blob data for the given transaction hash.
    fn get(&self, tx: B256) -> Result<Option<Arc<BlobTransactionSidecarVariant>>, BlobStoreError> {
        Ok(self.inner.store.read().get(&tx).cloned())
    }

    fn contains(&self, tx: B256) -> Result<bool, BlobStoreError> {
        Ok(self.inner.store.read().contains_key(&tx))
    }

    fn get_all(
        &self,
        txs: Vec<B256>,
    ) -> Result<Vec<(B256, Arc<BlobTransactionSidecarVariant>)>, BlobStoreError> {
        let store = self.inner.store.read();
        Ok(txs.into_iter().filter_map(|tx| store.get(&tx).map(|item| (tx, item.clone()))).collect())
    }

    fn get_exact(
        &self,
        txs: Vec<B256>,
    ) -> Result<Vec<Arc<BlobTransactionSidecarVariant>>, BlobStoreError> {
        if txs.is_empty() {
            return Ok(Vec::new());
        }
        let store = self.inner.store.read();
        txs.into_iter()
            .map(|tx| store.get(&tx).cloned().ok_or(BlobStoreError::MissingSidecar(tx)))
            .collect()
    }

    fn get_by_versioned_hashes_v1(
        &self,
        versioned_hashes: &[B256],
    ) -> Result<Vec<Option<BlobAndProofV1>>, BlobStoreError> {
        let mut result = vec![None; versioned_hashes.len()];
        for blob_sidecar in self.inner.store.read().values() {
            if let Some(blob_sidecar) = blob_sidecar.as_eip4844() {
                for (hash_idx, match_result) in
                    blob_sidecar.match_versioned_hashes(versioned_hashes)
                {
                    result[hash_idx] = Some(match_result);
                }
            }

            // Return early if all blobs are found.
            if result.iter().all(|blob| blob.is_some()) {
                break;
            }
        }
        Ok(result)
    }

    fn get_by_versioned_hashes_v2(
        &self,
        versioned_hashes: &[B256],
    ) -> Result<Option<Vec<BlobAndProofV2>>, BlobStoreError> {
        let result = self.get_by_versioned_hashes_eip7594(versioned_hashes);
        if result.iter().all(|blob| blob.is_some()) {
            Ok(Some(result.into_iter().map(Option::unwrap).collect()))
        } else {
            Ok(None)
        }
    }

    fn get_by_versioned_hashes_v3(
        &self,
        versioned_hashes: &[B256],
    ) -> Result<Vec<Option<BlobAndProofV2>>, BlobStoreError> {
        Ok(self.get_by_versioned_hashes_eip7594(versioned_hashes))
    }

    fn get_by_versioned_hashes_v4(
        &self,
        versioned_hashes: &[B256],
        indices_bitarray: B128,
    ) -> Result<Vec<Option<BlobCellsAndProofsV1>>, BlobStoreError> {
        self.get_by_versioned_hashes_cells_eip7594(versioned_hashes, indices_bitarray)
    }

    fn has_versioned_hashes(&self, versioned_hashes: &[B256]) -> Result<Vec<bool>, BlobStoreError> {
        let mut result = vec![false; versioned_hashes.len()];
        for blob_sidecar in self.inner.store.read().values() {
            for available_hash in blob_sidecar.versioned_hashes() {
                for (idx, requested_hash) in versioned_hashes.iter().enumerate() {
                    if !result[idx] && *requested_hash == available_hash {
                        result[idx] = true;
                    }
                }
            }

            if result.iter().all(|available| *available) {
                break;
            }
        }
        Ok(result)
    }

    fn get_cells(
        &self,
        tx: B256,
        indices_bitarray: B128,
    ) -> Result<Option<Vec<Cell>>, BlobStoreError> {
        let Some(sidecar) = self.get(tx)? else {
            return Ok(None);
        };

        let Some(sidecar) = sidecar.as_eip7594() else {
            return Ok(None);
        };

        sidecar
            .compute_matching_cells(BlobCellMask::new(indices_bitarray))
            .map(Some)
            .map_err(|err| BlobStoreError::Other(Box::new(err)))
    }

    fn data_size_hint(&self) -> Option<usize> {
        Some(self.inner.size_tracker.data_size())
    }

    fn blobs_len(&self) -> usize {
        self.inner.size_tracker.blobs_len()
    }
}

/// Removes the given blob from the store and returns the size of the blob that was removed.
#[inline]
fn remove_size(store: &mut B256Map<Arc<BlobTransactionSidecarVariant>>, tx: &B256) -> usize {
    store.remove(tx).map(|rem| rem.size()).unwrap_or_default()
}

/// Inserts the given blob into the store and returns the size of the blob that was added.
///
/// We don't need to handle the size updates for replacements because transactions are unique.
#[inline]
fn insert_size(
    store: &mut B256Map<Arc<BlobTransactionSidecarVariant>>,
    tx: B256,
    blob: BlobTransactionSidecarVariant,
) -> usize {
    let add = blob.size();
    store.insert(tx, Arc::new(blob));
    add
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy_consensus::BlobTransactionSidecar;
    use alloy_eips::{
        eip4844::{kzg_to_versioned_hash, Blob, BlobAndProofV2, Bytes48},
        eip7594::{
            BlobTransactionSidecarEip7594, BlobTransactionSidecarVariant, CELLS_PER_EXT_BLOB,
        },
    };

    fn eip7594_single_blob_sidecar() -> (BlobTransactionSidecarVariant, B256, BlobAndProofV2) {
        let blob = Blob::default();
        let commitment = Bytes48::default();
        let cell_proofs = vec![Bytes48::default(); CELLS_PER_EXT_BLOB];

        let versioned_hash = kzg_to_versioned_hash(commitment.as_slice());

        let expected =
            BlobAndProofV2 { blob: Box::new(Blob::default()), proofs: cell_proofs.clone() };
        let sidecar = BlobTransactionSidecarEip7594::new(vec![blob], vec![commitment], cell_proofs);

        (BlobTransactionSidecarVariant::Eip7594(sidecar), versioned_hash, expected)
    }

    fn eip4844_single_blob_sidecar() -> (BlobTransactionSidecarVariant, B256) {
        let blob = Blob::default();
        let commitment = Bytes48::from([1u8; 48]);
        let proof = Bytes48::default();
        let versioned_hash = kzg_to_versioned_hash(commitment.as_slice());
        let sidecar = BlobTransactionSidecar {
            blobs: vec![blob],
            commitments: vec![commitment],
            proofs: vec![proof],
        };

        (BlobTransactionSidecarVariant::Eip4844(sidecar), versioned_hash)
    }

    #[test]
    fn mem_has_blobs_returns_ordered_availability() {
        let store = InMemoryBlobStore::default();

        let (eip7594_sidecar, eip7594_hash, _) = eip7594_single_blob_sidecar();
        let (eip4844_sidecar, eip4844_hash) = eip4844_single_blob_sidecar();
        store.insert(B256::random(), eip7594_sidecar.into()).unwrap();
        store.insert(B256::random(), eip4844_sidecar.into()).unwrap();

        let request = vec![eip7594_hash, B256::ZERO, eip4844_hash, eip7594_hash];
        assert_eq!(store.has_versioned_hashes(&request).unwrap(), vec![true, false, true, true]);
    }

    #[test]
    fn mem_get_blobs_v3_returns_partial_results() {
        let store = InMemoryBlobStore::default();

        let (sidecar, versioned_hash, expected) = eip7594_single_blob_sidecar();
        store.insert(B256::random(), sidecar.into()).unwrap();

        assert_ne!(versioned_hash, B256::ZERO);

        let request = vec![versioned_hash, B256::ZERO];
        let v2 = store.get_by_versioned_hashes_v2(&request).unwrap();
        assert!(v2.is_none(), "v2 must return null if any requested blob is missing");

        let v3 = store.get_by_versioned_hashes_v3(&request).unwrap();
        assert_eq!(v3, vec![Some(expected), None]);
    }

    #[test]
    fn mem_get_blobs_v4_returns_requested_cells() {
        let store = InMemoryBlobStore::default();

        let (sidecar, versioned_hash, _) = eip7594_single_blob_sidecar();
        store.insert(B256::random(), sidecar.into()).unwrap();

        let indices_bitarray = B128::from((1u128 << 0) | (1u128 << 7));
        let request = vec![versioned_hash, B256::ZERO];

        let v4 = store.get_by_versioned_hashes_v4(&request, indices_bitarray).unwrap();
        assert_eq!(v4.len(), request.len());
        assert!(v4[1].is_none());

        let cells_and_proofs = v4[0].as_ref().unwrap();
        assert_eq!(cells_and_proofs.blob_cells.len(), 2);
        assert_eq!(cells_and_proofs.proofs.len(), 2);
        assert!(cells_and_proofs.blob_cells.iter().all(Option::is_some));
        assert_eq!(cells_and_proofs.proofs, vec![Some(Bytes48::default()); 2]);
    }

    #[test]
    fn mem_get_cells_returns_requested_cells() {
        let store = InMemoryBlobStore::default();

        let tx_hash = B256::random();
        let (sidecar, versioned_hash, _) = eip7594_single_blob_sidecar();
        store.insert(tx_hash, sidecar.into()).unwrap();

        let indices_bitarray = B128::from((1u128 << 0) | (1u128 << 7));
        let expected = store
            .get_by_versioned_hashes_v4(&[versioned_hash], indices_bitarray)
            .unwrap()
            .pop()
            .unwrap()
            .unwrap()
            .blob_cells
            .into_iter()
            .collect::<Option<Vec<_>>>()
            .unwrap();

        assert_eq!(store.get_cells(tx_hash, indices_bitarray).unwrap(), Some(expected));
    }
}
