use crate::AvailabilityProcessingStatus;
use crate::data_column_verification::KzgVerifiedCustodyDataColumn;
use crate::fetch_blobs::fetch_blobs_beacon_adapter::MockFetchBlobsBeaconAdapter;
use crate::fetch_blobs::{FetchEngineBlobError, fetch_and_process_engine_blobs_inner};
use crate::partial_data_column_assembler::PartialDataColumnAssembler;
use crate::test_utils::{EphemeralHarnessType, get_kzg};
use bls::Signature;
use eth2::types::BlobsBundle;
use execution_layer::json_structures::{BlobAndProof, BlobAndProofV1, BlobAndProofV2};
use execution_layer::test_utils::generate_blobs;
use maplit::hashset;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use task_executor::test_utils::TestRuntime;
use types::{
    BeaconBlock, BeaconBlockFulu, EmptyBlock, EthSpec, ForkName, Hash256, MainnetEthSpec,
    SignedBeaconBlock, SignedBeaconBlockFulu,
};

type E = MainnetEthSpec;
type T = EphemeralHarnessType<E>;

mod get_blobs_v2 {
    use super::*;
    use types::{ColumnIndex, PartialDataColumnHeader};

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_fetch_blobs_v2_no_blobs_in_block() {
        let mut mock_adapter = mock_beacon_adapter(ForkName::Fulu, false);
        let (publish_fn, _s) = mock_publish_fn();
        let block = SignedBeaconBlock::<E>::Fulu(SignedBeaconBlockFulu {
            message: BeaconBlockFulu::empty(mock_adapter.spec()),
            signature: Signature::empty(),
        });
        let block_root = block.canonical_root();

        // Expectations: engine fetch blobs should not be triggered
        mock_adapter.expect_get_blobs_v2().times(0);
        mock_adapter.expect_process_engine_blobs().times(0);

        let custody_columns: [ColumnIndex; 3] = [0, 1, 2];
        let processing_status = fetch_and_process_engine_blobs_inner(
            mock_adapter,
            block_root,
            Arc::new((&block).try_into().unwrap()),
            &custody_columns,
            publish_fn,
        )
        .await
        .expect("fetch blobs should succeed");

        assert_eq!(processing_status, None);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_fetch_blobs_v2_no_blobs_returned() {
        let mut mock_adapter = mock_beacon_adapter(ForkName::Fulu, false);
        let (publish_fn, _) = mock_publish_fn();
        let (block, _blobs_and_proofs) = create_test_block_and_blobs(&mock_adapter, 2);
        let block_root = block.canonical_root();

        // No blobs in EL response
        mock_get_blobs_v2_response(&mut mock_adapter, None);

        // Trigger fetch blobs on the block
        let custody_columns: [ColumnIndex; 3] = [0, 1, 2];
        let processing_status = fetch_and_process_engine_blobs_inner(
            mock_adapter,
            block_root,
            Arc::new(PartialDataColumnHeader::try_from(block.as_ref()).unwrap()),
            &custody_columns,
            publish_fn,
        )
        .await
        .expect("fetch blobs should succeed");

        assert_eq!(processing_status, None);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_fetch_blobs_v2_partial_blobs_returned() {
        let mut mock_adapter = mock_beacon_adapter(ForkName::Fulu, false);
        let (publish_fn, publish_fn_args) = mock_publish_fn();
        let (block, mut blobs_and_proofs) = create_test_block_and_blobs(&mock_adapter, 2);
        let block_root = block.canonical_root();

        // Missing blob in EL response
        blobs_and_proofs.pop();
        mock_get_blobs_v2_response(&mut mock_adapter, Some(blobs_and_proofs));
        // No blobs should be processed
        mock_adapter.expect_process_engine_blobs().times(0);

        // Trigger fetch blobs on the block
        let custody_columns: [ColumnIndex; 3] = [0, 1, 2];
        let processing_status = fetch_and_process_engine_blobs_inner(
            mock_adapter,
            block_root,
            Arc::new(PartialDataColumnHeader::try_from(block.as_ref()).unwrap()),
            &custody_columns,
            publish_fn,
        )
        .await
        .expect("fetch blobs should succeed");

        assert_eq!(processing_status, None);
        assert_eq!(
            publish_fn_args.lock().unwrap().len(),
            0,
            "no columns should be published"
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_fetch_blobs_v2_block_imported_after_el_response() {
        let mut mock_adapter = mock_beacon_adapter(ForkName::Fulu, false);
        let (publish_fn, publish_fn_args) = mock_publish_fn();
        let (block, blobs_and_proofs) = create_test_block_and_blobs(&mock_adapter, 2);
        let block_root = block.canonical_root();

        // All blobs returned, but fork choice already imported the block
        mock_get_blobs_v2_response(&mut mock_adapter, Some(blobs_and_proofs));
        mock_fork_choice_contains_block(&mut mock_adapter, vec![block.canonical_root()]);
        // No blobs should be processed
        mock_adapter.expect_process_engine_blobs().times(0);

        // Trigger fetch blobs on the block
        let custody_columns: [ColumnIndex; 3] = [0, 1, 2];
        let processing_status = fetch_and_process_engine_blobs_inner(
            mock_adapter,
            block_root,
            Arc::new(PartialDataColumnHeader::try_from(block.as_ref()).unwrap()),
            &custody_columns,
            publish_fn,
        )
        .await
        .expect("fetch blobs should succeed");

        assert_eq!(processing_status, None);
        assert_eq!(
            publish_fn_args.lock().unwrap().len(),
            0,
            "no columns should be published"
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_fetch_blobs_v2_no_new_columns_to_import() {
        let mut mock_adapter = mock_beacon_adapter(ForkName::Fulu, false);
        let (publish_fn, publish_fn_args) = mock_publish_fn();
        let (block, blobs_and_proofs) = create_test_block_and_blobs(&mock_adapter, 2);
        let block_root = block.canonical_root();

        // **GIVEN**:
        // All blobs returned
        mock_get_blobs_v2_response(&mut mock_adapter, Some(blobs_and_proofs));
        // block not yet imported into fork choice
        mock_fork_choice_contains_block(&mut mock_adapter, vec![]);
        // All data columns already seen on gossip
        mock_adapter
            .expect_data_column_known_for_observation_key()
            .returning(|_| Some(hashset![0, 1, 2]));
        // No blobs should be processed
        mock_adapter.expect_process_engine_blobs().times(0);

        // **WHEN**: Trigger `fetch_blobs` on the block
        let custody_columns: [ColumnIndex; 3] = [0, 1, 2];
        let processing_status = fetch_and_process_engine_blobs_inner(
            mock_adapter,
            block_root,
            Arc::new(PartialDataColumnHeader::try_from(block.as_ref()).unwrap()),
            &custody_columns,
            publish_fn,
        )
        .await
        .expect("fetch blobs should succeed");

        // **THEN**: Should NOT be processed and no columns should be published.
        assert_eq!(processing_status, None);
        assert_eq!(
            publish_fn_args.lock().unwrap().len(),
            0,
            "no columns should be published"
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_fetch_blobs_v2_success() {
        let mut mock_adapter = mock_beacon_adapter(ForkName::Fulu, false);
        let (publish_fn, publish_fn_args) = mock_publish_fn();
        let (block, blobs_and_proofs) = create_test_block_and_blobs(&mock_adapter, 2);
        let block_root = block.canonical_root();

        // All blobs returned, fork choice doesn't contain block
        mock_get_blobs_v2_response(&mut mock_adapter, Some(blobs_and_proofs));
        mock_fork_choice_contains_block(&mut mock_adapter, vec![]);
        mock_adapter
            .expect_data_column_known_for_observation_key()
            .returning(|_| None);
        mock_adapter
            .expect_cached_data_column_indexes()
            .returning(|_, _| None);
        mock_process_engine_blobs_result(
            &mut mock_adapter,
            Ok(AvailabilityProcessingStatus::Imported(block_root)),
        );

        // Trigger fetch blobs on the block
        let custody_columns: [ColumnIndex; 3] = [0, 1, 2];
        let processing_status = fetch_and_process_engine_blobs_inner(
            mock_adapter,
            block_root,
            Arc::new(PartialDataColumnHeader::try_from(block.as_ref()).unwrap()),
            &custody_columns,
            publish_fn,
        )
        .await
        .expect("fetch blobs should succeed");

        assert_eq!(
            processing_status,
            Some(AvailabilityProcessingStatus::Imported(block_root))
        );

        let published_columns = extract_published_blobs(publish_fn_args);
        assert!(
            matches!(
                published_columns,
                columns if columns.len() == custody_columns.len()
            ),
            "should publish custody columns"
        );
    }

    fn mock_get_blobs_v2_response(
        mock_adapter: &mut MockFetchBlobsBeaconAdapter<T>,
        blobs_and_proofs_opt: Option<Vec<BlobAndProof<E>>>,
    ) {
        let blobs_and_proofs_v2_opt = blobs_and_proofs_opt.map(|blobs_and_proofs| {
            blobs_and_proofs
                .into_iter()
                .map(|blob_and_proof| match blob_and_proof {
                    BlobAndProof::V2(inner) => inner,
                    _ => panic!("BlobAndProofV2 not expected"),
                })
                .collect()
        });
        mock_adapter
            .expect_get_blobs_v2()
            .return_once(move |_| Ok(blobs_and_proofs_v2_opt));
    }
}

/// Extract the `Vec<KzgVerifiedCustodyDataColumn<E>>` passed to the `publish_fn`.
fn extract_published_blobs(
    publish_fn_args: Arc<Mutex<Vec<Vec<KzgVerifiedCustodyDataColumn<E>>>>>,
) -> Vec<KzgVerifiedCustodyDataColumn<E>> {
    let mut calls = publish_fn_args.lock().unwrap();
    assert_eq!(calls.len(), 1);
    calls.pop().unwrap()
}

fn mock_process_engine_blobs_result(
    mock_adapter: &mut MockFetchBlobsBeaconAdapter<T>,
    result: Result<AvailabilityProcessingStatus, FetchEngineBlobError>,
) {
    mock_adapter
        .expect_process_engine_blobs()
        .return_once(move |_, _, _| result);
}

fn mock_fork_choice_contains_block(
    mock_adapter: &mut MockFetchBlobsBeaconAdapter<T>,
    block_roots: Vec<Hash256>,
) {
    mock_adapter
        .expect_fork_choice_contains_block()
        .returning(move |block_root| block_roots.contains(block_root));
}

fn create_test_block_and_blobs(
    mock_adapter: &MockFetchBlobsBeaconAdapter<T>,
    blob_count: usize,
) -> (Arc<SignedBeaconBlock<E>>, Vec<BlobAndProof<E>>) {
    let mut block =
        SignedBeaconBlock::from_block(BeaconBlock::empty(mock_adapter.spec()), Signature::empty());
    let fork = block.fork_name_unchecked();
    let (blobs_bundle, _tx) = generate_blobs::<E>(blob_count, fork).unwrap();
    let BlobsBundle {
        commitments,
        proofs,
        blobs,
    } = blobs_bundle;

    *block
        .message_mut()
        .body_mut()
        .blob_kzg_commitments_mut()
        .unwrap() = commitments;

    let blobs_and_proofs = if fork.fulu_enabled() {
        let proofs_len = proofs.len() / blobs.len();
        blobs
            .into_iter()
            .zip(proofs.chunks(proofs_len))
            .map(|(blob, proofs)| {
                BlobAndProof::V2(BlobAndProofV2 {
                    blob,
                    proofs: proofs.to_vec().try_into().unwrap(),
                })
            })
            .collect()
    } else {
        blobs
            .into_iter()
            .zip(proofs)
            .map(|(blob, proof)| BlobAndProof::V1(BlobAndProofV1 { blob, proof }))
            .collect()
    };

    (Arc::new(block), blobs_and_proofs)
}

#[allow(clippy::type_complexity)]
fn mock_publish_fn() -> (
    impl Fn(Vec<KzgVerifiedCustodyDataColumn<E>>) + Send + 'static,
    Arc<Mutex<Vec<Vec<KzgVerifiedCustodyDataColumn<E>>>>>,
) {
    // Keep track of the arguments captured by `publish_fn`.
    let captured_args = Arc::new(Mutex::new(vec![]));
    let captured_args_clone = captured_args.clone();
    let publish_fn = move |args| {
        let mut lock = captured_args_clone.lock().unwrap();
        lock.push(args);
    };
    (publish_fn, captured_args)
}

fn mock_beacon_adapter(fork_name: ForkName, get_blobs_v3: bool) -> MockFetchBlobsBeaconAdapter<T> {
    let test_runtime = TestRuntime::default();
    let spec = Arc::new(fork_name.make_genesis_spec(E::default_spec()));
    let kzg = get_kzg(&spec);
    let partial_assembler = PartialDataColumnAssembler::new(NonZeroUsize::new(32).unwrap());

    let mut mock_adapter = MockFetchBlobsBeaconAdapter::default();
    mock_adapter.expect_spec().return_const(spec.clone());
    mock_adapter.expect_kzg().return_const(kzg.clone());
    mock_adapter
        .expect_executor()
        .return_const(test_runtime.task_executor.clone());
    mock_adapter
        .expect_supports_get_blobs_v3()
        .returning(move || Ok(get_blobs_v3));
    mock_adapter
        .expect_partial_assembler()
        .return_const(Some(Arc::new(partial_assembler)));
    mock_adapter
}
