use crate::sandbox::{
    common::{
        call_contract_key_generation, execute_key_generation_and_add_random_state, gen_accounts,
        init, propose_and_vote_contract_binary, submit_attestations,
    },
    utils::{
        consts::PARTICIPANT_LEN,
        contract_build::current_contract,
        mpc_contract::{get_participants, get_state, get_tee_accounts},
        shared_key_utils::DomainKey,
        sign_utils::{make_and_submit_requests, submit_ckd_response, submit_signature_response},
    },
};
use mpc_contract::primitives::{
    key_state::{EpochId, Keyset},
    participants::Participants,
    thresholds::{Threshold, ThresholdParameters},
};
use near_account_id::AccountId;
use near_mpc_bounded_collections::NonEmptyBTreeSet;
use near_mpc_contract_interface::method_names;
use near_mpc_contract_interface::types as dtos;
use near_mpc_contract_interface::types::ProtocolContractState;
use near_mpc_contract_interface::types::{
    CKDResponse, DomainConfig, DomainPurpose, Protocol, ReconstructionThreshold,
};
use near_mpc_sdk::sign::SignatureRequestResponse;
use near_workspaces::{Account, Contract, Worker, network::Sandbox};
use rand_core::OsRng;
use rstest::rstest;
use std::collections::HashSet;
use std::collections::{BTreeMap, BTreeSet};

#[derive(Debug, Clone, Copy)]
enum Network {
    Testnet,
    Mainnet,
}

fn contract_code(network: Network) -> &'static [u8] {
    match network {
        Network::Mainnet => contract_history::current_mainnet(),
        Network::Testnet => contract_history::current_testnet(),
    }
}

async fn init_old_contract(
    worker: &Worker<Sandbox>,
    contract: &Contract,
    number_of_participants: usize,
) -> anyhow::Result<(Vec<Account>, Participants)> {
    let (accounts, participants) = gen_accounts(worker, number_of_participants).await;

    let threshold = ((participants.len() as f64) * 0.6).ceil() as u64;
    let threshold = Threshold::new(threshold);
    let threshold_parameters = ThresholdParameters::new(participants.clone(), threshold).unwrap();

    contract
        .call(method_names::INIT)
        .args_json(serde_json::json!({
            "parameters": &threshold_parameters,
        }))
        .transact()
        .await?
        .into_result()?;
    Ok((accounts, participants))
}

async fn healthcheck(contract: &Contract) -> anyhow::Result<bool> {
    let status = contract
        .call(method_names::STATE)
        .transact()
        .await?
        .into_result()
        .is_ok();
    Ok(status)
}

async fn deploy_old(worker: &Worker<Sandbox>, network: Network) -> anyhow::Result<Contract> {
    let old_wasm = contract_code(network);
    let old_contract = worker.dev_deploy(old_wasm).await?;
    Ok(old_contract)
}

async fn upgrade_to_new(old_contract: Contract) -> anyhow::Result<Contract> {
    let new_wasm = current_contract();
    let new_contract = old_contract
        .as_account()
        .deploy(new_wasm)
        .await?
        .into_result()?;
    Ok(new_contract)
}

/// Migrates the contract to a current contract build
/// and sanity checks that the upgraded code matches compiled contract bytes.
async fn migrate_and_assert_contract_code(contract: &Contract) -> anyhow::Result<()> {
    contract
        .call(method_names::MIGRATE)
        .transact()
        .await?
        .into_result()?;
    let code_hash_post_upgrade = contract.view_code().await.unwrap();
    let current_code_hash = current_contract();

    assert_eq!(*current_code_hash, code_hash_post_upgrade);

    Ok(())
}

/// Checks the contract in the following order:
/// 1. Are there any state-breaking changes?
/// 2. If so, does `migrate()` still work correctly?
///
/// These checks use the previous contract version (the one that introduced breaking changes)
/// as a baseline. If step 2 fails, you will be prompted to update the baseline contract.
#[rstest]
#[tokio::test]
async fn back_compatibility_without_state(
    #[values(Network::Mainnet, Network::Testnet)] network: Network,
) -> anyhow::Result<()> {
    let worker = near_workspaces::sandbox_with_version(test_utils::DEFAULT_SANDBOX_VERSION).await?;

    let contract = deploy_old(&worker, network).await?;

    init_old_contract(&worker, &contract, PARTICIPANT_LEN).await?;

    assert!(healthcheck(&contract).await?);

    let contract = upgrade_to_new(contract).await?;

    if healthcheck(&contract).await? {
        println!("✅ Back compatibility check succeeded: no breaking changes found 🫧.");
        return Ok(());
    }

    println!("🟨 Found breaking changes in the contract state.");
    println!("⚙️ Testing migration() call...");

    migrate_and_assert_contract_code(&contract)
        .await
        .expect("❌ Back compatibility check failed: migration() failed");

    if healthcheck(&contract).await? {
        println!("✅ Back compatibility check succeeded: migration() works fine 👍");
        return Ok(());
    };

    anyhow::bail!(
        "❌Back compatibility check failed: state() call doesnt work after migration(). Probably you should introduce new logic to the `migrate()` method."
    )
}

/// Ensures that contracts deployed with the production binary (Mainnet or Testnet)
/// can be upgraded to the [`current_contract`] binary using the proposal-and-vote flow.
#[rstest]
#[tokio::test]
async fn propose_upgrade_from_production_to_current_binary(
    #[values(Network::Mainnet, Network::Testnet)] network: Network,
) {
    let worker = near_workspaces::sandbox_with_version(test_utils::DEFAULT_SANDBOX_VERSION)
        .await
        .unwrap();
    let contract = deploy_old(&worker, network).await.unwrap();
    let (accounts, participants) = init_old_contract(&worker, &contract, PARTICIPANT_LEN)
        .await
        .unwrap();

    submit_attestations(&contract, &accounts, &participants).await;

    // Add state so migration logic is exercised
    execute_key_generation_and_add_random_state(
        &accounts,
        participants,
        &contract,
        &worker,
        &mut OsRng,
    )
    .await;

    let state_pre_upgrade: ProtocolContractState = get_state(&contract).await;

    propose_and_vote_contract_binary(&accounts, &contract, current_contract()).await;

    let state_post_upgrade: ProtocolContractState = get_state(&contract).await;

    assert_eq!(
        state_pre_upgrade, state_post_upgrade,
        "State of the contract should remain the same post upgrade."
    );
}

//// Verifies that upgrading the contract preserves state and functionality.
///
/// This test:
/// 1. Deploys an older version of the contract.
/// 2. Initializes it with participants and submits a parameter update proposal.
/// 3. Adds multiple domains with both `Ed25519` and `Secp256k1` schemes.
/// 4. Submits pending signature requests across those domains.
/// 5. Captures the full pre-upgrade state.
/// 6. Upgrades the contract to the new version and runs `migrate()`.
/// 7. Asserts that the state (participants, domains, proposals, signature requests, etc.)
///    is identical post-upgrade.
/// 8. Confirms that pending signature requests created before the upgrade
///    can still be responded to afterward.
#[rstest]
#[tokio::test]
async fn upgrade_preserves_state_and_requests(
    #[values(Network::Mainnet, Network::Testnet)] network: Network,
) {
    let worker = near_workspaces::sandbox_with_version(test_utils::DEFAULT_SANDBOX_VERSION)
        .await
        .unwrap();
    let contract = deploy_old(&worker, network).await.unwrap();
    let (accounts, participants) = init_old_contract(&worker, &contract, PARTICIPANT_LEN)
        .await
        .unwrap();

    let attested_account = &accounts[0];

    submit_attestations(&contract, &accounts, &participants).await;

    let injected_contract_state = execute_key_generation_and_add_random_state(
        &accounts,
        participants,
        &contract,
        &worker,
        &mut OsRng,
    )
    .await;

    let state_pre_upgrade: ProtocolContractState = get_state(&contract).await;

    assert!(healthcheck(&contract).await.unwrap());
    let contract = upgrade_to_new(contract).await.unwrap();
    migrate_and_assert_contract_code(&contract)
        .await
        .expect("❌ migration() failed");

    let state_post_upgrade: ProtocolContractState = get_state(&contract).await;

    assert_eq!(
        state_pre_upgrade, state_post_upgrade,
        "State of the contract should remain the same post upgrade."
    );

    for pending in injected_contract_state.pending_sign_requests {
        submit_signature_response(&pending.response, &contract, attested_account)
            .await
            .unwrap();

        let execution = pending.transaction.await.unwrap().into_result().unwrap();
        let returned: SignatureRequestResponse = execution.json().unwrap();

        assert_eq!(
            returned, pending.response.response,
            "Returned signature response does not match"
        );
    }
}

/// During the soft-launch transition every participant re-submits their TEE
/// attestation on the old contract (populating the previously-optional
/// `account_public_key` field). The #1710 migration drops any stored entry
/// that still has a missing account key, so this test reproduces the
/// production sequence: soft-launch re-submissions first, then the upgrade,
/// and verifies that every initial participant still has a stored attestation
/// after migration.
#[tokio::test]
async fn all_participants_get_valid_mock_attestation_for_soft_launch_upgrade() -> anyhow::Result<()>
{
    let worker = near_workspaces::sandbox_with_version(test_utils::DEFAULT_SANDBOX_VERSION).await?;
    let contract = deploy_old(&worker, Network::Testnet).await?;

    let (accounts, participants) = init_old_contract(&worker, &contract, PARTICIPANT_LEN).await?;

    let initial_participants = get_participants(&contract).await?;
    let participant_set_is_not_empty = !initial_participants.participants.is_empty();
    assert!(
        participant_set_is_not_empty,
        "Test must contain a contract with at least one participant"
    );

    submit_attestations(&contract, &accounts, &participants).await;

    let contract = upgrade_to_new(contract).await?;

    migrate_and_assert_contract_code(&contract)
        .await
        .expect("❌ Back compatibility check failed: migration() failed");

    let accounts_with_tee_attestation_post_upgrade: HashSet<AccountId> =
        get_tee_accounts(&contract)
            .await
            .unwrap()
            .into_iter()
            .map(|node_id| node_id.account_id.clone())
            .collect();

    let participant_set: HashSet<AccountId> = initial_participants
        .participants
        .iter()
        .map(|(account_id, _, _)| account_id.clone())
        .collect();

    assert_eq!(
        accounts_with_tee_attestation_post_upgrade, participant_set,
        "All initial participants must have a valid attestation post upgrade."
    );
    Ok(())
}

//// Verifies that upgrading the contract preserves state and allows the new
/// functionality, in this case only CKD
///
/// This test:
/// 1. Deploys an older version of the contract.
/// 2. Initializes it with participants and submits a parameter update proposal.
/// 3. Adds multiple domains with both `Ed25519` and `Secp256k1` schemes.
/// 4. Submits pending signature requests across those domains.
/// 5. Captures the full pre-upgrade state.
/// 6. Upgrades the contract to the new version and runs `migrate()`.
/// 7. Asserts that the state (participants, domains, proposals, signature requests, etc.)
///    is identical post-upgrade.
/// 10. Adds new domains, including CKD
/// 11. Submits new signature and ckd requests
/// 12. Confirms that pending signature and ckd requests created before and after the upgrade
///    can still be responded to.
#[rstest]
#[tokio::test]
async fn upgrade_allows_new_request_types(
    #[values(Network::Mainnet, Network::Testnet)] network: Network,
) {
    let rng = &mut OsRng;

    let worker = near_workspaces::sandbox_with_version(test_utils::DEFAULT_SANDBOX_VERSION)
        .await
        .unwrap();
    let contract = deploy_old(&worker, network).await.unwrap();
    let (accounts, participants) = init_old_contract(&worker, &contract, PARTICIPANT_LEN)
        .await
        .unwrap();
    let attested_account = &accounts[0];

    submit_attestations(&contract, &accounts, &participants).await;

    let injected_contract_state = execute_key_generation_and_add_random_state(
        &accounts,
        participants,
        &contract,
        &worker,
        rng,
    )
    .await;

    let state_pre_upgrade: ProtocolContractState = get_state(&contract).await;

    assert!(healthcheck(&contract).await.unwrap());
    let contract = upgrade_to_new(contract).await.unwrap();
    migrate_and_assert_contract_code(&contract)
        .await
        .expect("❌ migration() failed");

    let state_post_upgrade: ProtocolContractState = get_state(&contract).await;

    assert_eq!(
        state_pre_upgrade, state_post_upgrade,
        "State of the contract should remain the same post upgrade."
    );

    let first_available_domain_id = injected_contract_state.domain_keys.len() as u64;

    // 2. Add new domains
    let domains_to_add = [
        DomainConfig {
            id: first_available_domain_id.into(),
            protocol: Protocol::ConfidentialKeyDerivation,
            reconstruction_threshold: ReconstructionThreshold::new(6),
            purpose: DomainPurpose::CKD,
        },
        DomainConfig {
            id: (first_available_domain_id + 1).into(),
            protocol: Protocol::Frost,
            reconstruction_threshold: ReconstructionThreshold::new(6),
            purpose: DomainPurpose::Sign,
        },
    ];

    const EPOCH_ID: u64 = 0;
    let added_domain_keys =
        call_contract_key_generation(&domains_to_add, &accounts, &contract, EPOCH_ID).await;

    let current_keys: Vec<DomainKey> = injected_contract_state
        .domain_keys
        .clone()
        .iter()
        .chain(added_domain_keys.iter())
        .cloned()
        .collect();

    let (pending_sign_requests, pending_ckd_requests) =
        make_and_submit_requests(&current_keys, &contract, &worker, rng).await;

    for pending in injected_contract_state
        .pending_sign_requests
        .into_iter()
        .chain(pending_sign_requests.into_iter())
    {
        submit_signature_response(&pending.response, &contract, attested_account)
            .await
            .unwrap();

        let execution = pending.transaction.await.unwrap().into_result().unwrap();
        let returned: SignatureRequestResponse = execution.json().unwrap();

        assert_eq!(
            returned, pending.response.response,
            "Returned signature response does not match"
        );
    }

    for pending in pending_ckd_requests {
        submit_ckd_response(&pending.ckd_response, &contract, attested_account)
            .await
            .unwrap();

        let execution = pending.transaction.await.unwrap().into_result().unwrap();
        let returned: CKDResponse = execution.json().unwrap();

        assert_eq!(
            returned, pending.ckd_response.response,
            "Returned ckd response does not match"
        );
    }
}

#[tokio::test]
async fn init_running_rejects_external_callers_pre_initialization() {
    let (worker, contract) = init().await;
    let number_of_participants = 2;
    let (accounts, participants) = gen_accounts(&worker, number_of_participants).await;

    let threshold_parameters = ThresholdParameters::new(
        participants.clone(),
        Threshold::new(number_of_participants as u64),
    )
    .unwrap();

    let init_running_args = serde_json::json!({
            "domains": [],
            "next_domain_id": 0,
            "keyset": Keyset::new(EpochId::new(2), vec![]),
            "parameters": threshold_parameters,
    });

    let execution_error = accounts[0]
        .call(contract.id(), method_names::INIT_RUNNING)
        .max_gas()
        .args_json(init_running_args)
        .transact()
        .await
        .unwrap()
        .into_result()
        .expect_err("method is private and not callable from participant account.");

    let error_message = format!("{:?}", execution_error);

    let expected_error_message = "Smart contract panicked: Method init_running is private";

    assert!(
        error_message.contains(expected_error_message),
        "init_running call was accepted by external caller. expected method to be private. {:?}",
        error_message
    )
}

/// Verifies that per-node foreign chain configurations registered on the old
/// contract via the deprecated `register_foreign_chain_config` are migrated to
/// the new `node_foreign_chain_support` layout: each node's full
/// `ForeignChainConfiguration` (chain → RPC providers) collapses to the set of
/// supported chains, and per-node entries are preserved (not merged).
#[rstest]
#[tokio::test]
async fn upgrade_preserves_per_node_foreign_chain_support(
    #[values(Network::Mainnet, Network::Testnet)] network: Network,
) -> anyhow::Result<()> {
    // Three participants, each registering a distinct chain configuration. The
    // chosen sets are deliberately overlapping but not equal so the test can
    // detect any per-node merging or loss.
    let per_node_chains: [&[dtos::ForeignChain]; 3] = [
        &[dtos::ForeignChain::Bitcoin, dtos::ForeignChain::Ethereum],
        &[dtos::ForeignChain::Bitcoin],
        &[dtos::ForeignChain::Solana],
    ];

    // Given: an old contract with participants and per-node foreign chain
    // configurations registered through the deprecated method.
    let worker = near_workspaces::sandbox_with_version(test_utils::DEFAULT_SANDBOX_VERSION).await?;
    let contract = deploy_old(&worker, network).await?;
    let (accounts, _participants) =
        init_old_contract(&worker, &contract, per_node_chains.len()).await?;

    for (account, chains) in accounts.iter().zip(per_node_chains.iter()) {
        #[expect(deprecated)]
        let configuration: dtos::ForeignChainConfiguration = chains
            .iter()
            .map(|chain| {
                (
                    *chain,
                    NonEmptyBTreeSet::new(dtos::RpcProvider {
                        rpc_url: format!("https://{:?}.{}.example.near", chain, account.id()),
                    }),
                )
            })
            .collect::<BTreeMap<_, _>>()
            .into();

        #[expect(deprecated)]
        account
            .call(contract.id(), method_names::REGISTER_FOREIGN_CHAIN_CONFIG)
            .args_json(serde_json::json!({
                "foreign_chain_configuration": configuration,
            }))
            .transact()
            .await?
            .into_result()?;
    }

    // When: we upgrade the contract and run migrate.
    let contract = upgrade_to_new(contract).await?;
    migrate_and_assert_contract_code(&contract)
        .await
        .expect("❌ migration() failed");

    // Then: each node's supported-chain set matches the chains it originally
    // registered (RPC providers are dropped by the new layout).
    let support: dtos::ForeignChainSupportByNode = contract
        .view(method_names::GET_FOREIGN_CHAIN_SUPPORT_BY_NODE)
        .await?
        .json()?;

    for (account, chains) in accounts.iter().zip(per_node_chains.iter()) {
        let actual = support
            .foreign_chain_support_by_node
            .get(account.id())
            .unwrap_or_else(|| panic!("entry for {} preserved post-upgrade", account.id()));
        let expected: dtos::SupportedForeignChains =
            chains.iter().copied().collect::<BTreeSet<_>>().into();
        assert_eq!(
            *actual,
            expected,
            "supported chains for {} should match what was registered pre-upgrade",
            account.id(),
        );
    }
    Ok(())
}