use crate::data_column_verification::KzgVerifiedCustodyDataColumn;
use crate::fetch_blobs::FetchEngineBlobError;
use crate::observed_data_sidecars::ObservationKey;
use crate::partial_data_column_assembler::PartialDataColumnAssembler;
use crate::{AvailabilityProcessingStatus, BeaconChain, BeaconChainTypes};
use execution_layer::json_structures::{BlobAndProofV2, BlobAndProofV3};
use kzg::Kzg;
#[cfg(test)]
use mockall::automock;
use std::collections::HashSet;
use std::sync::Arc;
use task_executor::TaskExecutor;
use types::{ChainSpec, ColumnIndex, Hash256, Slot};

/// An adapter to the `BeaconChain` functionalities to remove `BeaconChain` from direct dependency to enable testing fetch blobs logic.
pub(crate) struct FetchBlobsBeaconAdapter<T: BeaconChainTypes> {
    chain: Arc<BeaconChain<T>>,
    spec: Arc<ChainSpec>,
}

#[cfg_attr(test, automock, allow(dead_code))]
impl<T: BeaconChainTypes> FetchBlobsBeaconAdapter<T> {
    pub(crate) fn new(chain: Arc<BeaconChain<T>>) -> Self {
        let spec = chain.spec.clone();
        Self { chain, spec }
    }

    pub(crate) fn spec(&self) -> &Arc<ChainSpec> {
        &self.spec
    }

    pub(crate) fn kzg(&self) -> &Arc<Kzg> {
        &self.chain.kzg
    }

    pub(crate) fn executor(&self) -> &TaskExecutor {
        &self.chain.task_executor
    }

    pub(crate) fn partial_assembler(&self) -> Option<Arc<PartialDataColumnAssembler<T::EthSpec>>> {
        self.chain
            .data_availability_checker
            .partial_assembler()
            .cloned()
    }

    pub(crate) async fn get_blobs_v2(
        &self,
        versioned_hashes: Vec<Hash256>,
    ) -> Result<Option<Vec<BlobAndProofV2<T::EthSpec>>>, FetchEngineBlobError> {
        let execution_layer = self
            .chain
            .execution_layer
            .as_ref()
            .ok_or(FetchEngineBlobError::ExecutionLayerMissing)?;

        execution_layer
            .get_blobs_v2(versioned_hashes)
            .await
            .map_err(FetchEngineBlobError::RequestFailed)
    }

    pub(crate) async fn get_blobs_v3(
        &self,
        versioned_hashes: Vec<Hash256>,
    ) -> Result<Option<Vec<BlobAndProofV3<T::EthSpec>>>, FetchEngineBlobError> {
        let execution_layer = self
            .chain
            .execution_layer
            .as_ref()
            .ok_or(FetchEngineBlobError::ExecutionLayerMissing)?;

        execution_layer
            .get_blobs_v3(versioned_hashes)
            .await
            .map_err(FetchEngineBlobError::RequestFailed)
    }

    pub(crate) fn data_column_known_for_observation_key(
        &self,
        observation_key: ObservationKey,
    ) -> Option<HashSet<ColumnIndex>> {
        self.chain
            .observed_column_sidecars
            .read()
            .known_for_observation_key(&observation_key)
            .cloned()
    }

    pub(crate) fn cached_data_column_indexes(
        &self,
        block_root: &Hash256,
        slot: Slot,
    ) -> Option<Vec<u64>> {
        self.chain.cached_data_column_indexes(block_root, slot)
    }

    pub(crate) async fn process_engine_blobs(
        &self,
        slot: Slot,
        block_root: Hash256,
        blobs: Vec<KzgVerifiedCustodyDataColumn<T::EthSpec>>,
    ) -> Result<AvailabilityProcessingStatus, FetchEngineBlobError> {
        self.chain
            .process_engine_blobs(slot, block_root, blobs)
            .await
            .map_err(FetchEngineBlobError::BlobProcessingError)
    }

    pub(crate) fn fork_choice_contains_block(&self, block_root: &Hash256) -> bool {
        self.chain
            .canonical_head
            .fork_choice_read_lock()
            .contains_block(block_root)
    }

    pub(crate) async fn supports_get_blobs_v3(&self) -> Result<bool, FetchEngineBlobError> {
        let execution_layer = self
            .chain
            .execution_layer
            .as_ref()
            .ok_or(FetchEngineBlobError::ExecutionLayerMissing)?;

        execution_layer
            .get_engine_capabilities(None)
            .await
            .map_err(FetchEngineBlobError::RequestFailed)
            .map(|caps| caps.get_blobs_v3)
    }
}
