use crate::indexer::participants::ContractState;
use crate::p2p::testing::PortSeed;
use crate::tests::{
    DEFAULT_BLOCK_TIME, DEFAULT_MAX_PROTOCOL_WAIT_TIME, DEFAULT_MAX_SIGNATURE_WAIT_TIME,
    IntegrationTestSetup, request_ckd_and_await_response, request_ckd_pv_and_await_response,
    request_signature_and_await_response,
};
use crate::tracking::AutoAbortTask;
use mpc_primitives::domain::DomainId;
use near_mpc_contract_interface::types::{
    DomainConfig, DomainPurpose, Protocol, ReconstructionThreshold,
};
use near_time::Clock;

// Make a cluster of four nodes, test that we can generate keyshares
// and then produce signatures.
#[tokio::test]
#[test_log::test]
async fn test_basic_cluster() {
    const NUM_PARTICIPANTS: usize = 4;
    const THRESHOLD: usize = 3;
    const TXN_DELAY_BLOCKS: u64 = 1;
    let temp_dir = tempfile::tempdir().unwrap();
    let mut setup: IntegrationTestSetup = IntegrationTestSetup::new(
        Clock::real(),
        temp_dir.path(),
        (0..NUM_PARTICIPANTS)
            .map(|i| format!("test{}", i).parse().unwrap())
            .collect(),
        THRESHOLD,
        TXN_DELAY_BLOCKS,
        PortSeed::BASIC_CLUSTER_TEST,
        DEFAULT_BLOCK_TIME,
    );

    let signature_domain_ecdsa = DomainConfig {
        id: DomainId(0),
        protocol: Protocol::CaitSith,
        reconstruction_threshold: ReconstructionThreshold::new(3),
        purpose: DomainPurpose::Sign,
    };

    let signature_domain_eddsa = DomainConfig {
        id: DomainId(1),
        protocol: Protocol::Frost,
        reconstruction_threshold: ReconstructionThreshold::new(3),
        purpose: DomainPurpose::Sign,
    };

    let ckd_domain = DomainConfig {
        id: DomainId(2),
        protocol: Protocol::ConfidentialKeyDerivation,
        reconstruction_threshold: ReconstructionThreshold::new(3),
        purpose: DomainPurpose::CKD,
    };

    let domains = vec![
        signature_domain_ecdsa.clone(),
        signature_domain_eddsa.clone(),
        ckd_domain.clone(),
    ];

    {
        let mut contract = setup.indexer.contract_mut().await;
        contract.initialize(setup.participants.clone());
        contract.add_domains(domains.clone());
    }

    let _runs = setup
        .configs
        .into_iter()
        .map(|config| AutoAbortTask::from(tokio::spawn(config.run())))
        .collect::<Vec<_>>();

    setup
        .indexer
        .wait_for_contract_state(
            |state| matches!(state, ContractState::Running(_)),
            DEFAULT_MAX_PROTOCOL_WAIT_TIME * domains.len() as u32,
        )
        .await
        .expect("timeout waiting for keygen to complete");

    assert!(
        request_signature_and_await_response(
            &mut setup.indexer,
            "user0",
            &signature_domain_ecdsa,
            DEFAULT_MAX_SIGNATURE_WAIT_TIME
        )
        .await
        .is_some()
    );

    assert!(
        request_signature_and_await_response(
            &mut setup.indexer,
            "user0",
            &signature_domain_eddsa,
            DEFAULT_MAX_SIGNATURE_WAIT_TIME
        )
        .await
        .is_some()
    );

    assert!(
        request_ckd_and_await_response(
            &mut setup.indexer,
            "user0",
            &ckd_domain,
            DEFAULT_MAX_SIGNATURE_WAIT_TIME
        )
        .await
        .is_some()
    );

    assert!(
        request_ckd_pv_and_await_response(
            &mut setup.indexer,
            "user0",
            &ckd_domain,
            DEFAULT_MAX_SIGNATURE_WAIT_TIME
        )
        .await
        .is_some()
    );
}