use crate::sandbox::utils::{
    consts::{GAS_FOR_VOTE_CANCEL_RESHARING, GAS_FOR_VOTE_NEW_PARAMETERS, GAS_FOR_VOTE_RESHARED},
    mpc_contract::get_state,
    transactions::execute_async_transactions,
};
use dtos::{AttemptId, EpochId, KeyEventId};
use mpc_contract::primitives::thresholds::{ProposedThresholdParameters, ThresholdParameters};
use near_mpc_contract_interface::method_names;
use near_mpc_contract_interface::types::{self as dtos, ProtocolContractState};
use near_workspaces::{Account, Contract};
use serde_json::json;
use std::collections::BTreeMap;

pub async fn conclude_resharing(
    contract: &Contract,
    all_participants: &[Account],
    prospective_epoch_id: EpochId,
) -> anyhow::Result<()> {
    let ProtocolContractState::Resharing(resharing_state) = get_state(contract).await else {
        anyhow::bail!("expected resharing state");
    };
    if resharing_state.resharing_key.epoch_id != prospective_epoch_id {
        anyhow::bail!("epoch id mismatch");
    }
    let domain_configs = resharing_state
        .previous_running_state
        .domains
        .domains
        .clone();
    for domain_config in &domain_configs {
        let key_event_id = KeyEventId {
            epoch_id: prospective_epoch_id,
            domain_id: domain_config.id,
            attempt_id: AttemptId(0),
        };
        let state = get_state(contract).await;
        if !matches!(state, ProtocolContractState::Resharing(_)) {
            anyhow::bail!("expected resharing state");
        }
        start_reshare_instance(contract, all_participants, key_event_id).await?;

        vote_reshared(contract, all_participants, key_event_id).await?;
    }
    Ok(())
}

pub async fn vote_cancel_reshaing(contract: &Contract, accounts: &[Account]) -> anyhow::Result<()> {
    execute_async_transactions(
        accounts,
        contract,
        method_names::VOTE_CANCEL_RESHARING,
        &json!({}),
        GAS_FOR_VOTE_CANCEL_RESHARING,
    )
    .await
}

pub async fn vote_new_parameters(
    contract: &Contract,
    prospective_epoch_id: u64,
    proposal: &ThresholdParameters,
    persistent_participants: &[Account],
    new_participants: &[Account],
) -> anyhow::Result<()> {
    let proposal = ProposedThresholdParameters::new(proposal.clone(), BTreeMap::new());
    let json_args = json!({
        "prospective_epoch_id": prospective_epoch_id,
        "proposal": proposal,
    });
    // At least threshold old participants need to vote first,
    // here we are just using all of them
    execute_async_transactions(
        persistent_participants,
        contract,
        method_names::VOTE_NEW_PARAMETERS,
        &json_args,
        GAS_FOR_VOTE_NEW_PARAMETERS,
    )
    .await?;

    // then new participant can vote
    execute_async_transactions(
        new_participants,
        contract,
        method_names::VOTE_NEW_PARAMETERS,
        &json_args,
        GAS_FOR_VOTE_NEW_PARAMETERS,
    )
    .await?;
    Ok(())
}

pub async fn start_reshare_instance(
    contract: &Contract,
    accounts: &[Account],
    key_event_id: KeyEventId,
) -> anyhow::Result<()> {
    let state = get_state(contract).await;
    let active = match &state {
        ProtocolContractState::Initializing(s) => {
            &s.generating_key.parameters.participants.participants
        }
        ProtocolContractState::Running(s) => &s.parameters.participants.participants,
        ProtocolContractState::Resharing(s) => {
            &s.resharing_key.parameters.participants.participants
        }
        ProtocolContractState::NotInitialized => {
            panic!("protocol state must be initialized")
        }
    };
    let leader = accounts
        .iter()
        .min_by_key(|a| {
            active
                .iter()
                .find(|(account_id, _, _)| account_id == a.id())
                .map(|(_, pid, _)| *pid)
                .unwrap()
        })
        .unwrap();
    let result = leader
        .call(contract.id(), method_names::START_RESHARE_INSTANCE)
        .args_json(json!({"key_event_id": key_event_id}))
        .transact()
        .await?;
    if !result.is_success() {
        anyhow::bail!("{result:#?}");
    }
    Ok(())
}

pub async fn vote_reshared(
    contract: &Contract,
    accounts: &[Account],
    key_event_id: KeyEventId,
) -> anyhow::Result<()> {
    execute_async_transactions(
        accounts,
        contract,
        method_names::VOTE_RESHARED,
        &json!({"key_event_id": key_event_id}),
        GAS_FOR_VOTE_RESHARED,
    )
    .await
}

/// Performs a complete resharing operation with the given parameters.
/// This includes voting for new parameters, starting reshare instances for each domain,
/// and voting reshared to complete the transition.
pub async fn do_resharing(
    remaining_accounts: &[Account],
    contract: &Contract,
    new_threshold_parameters: ThresholdParameters,
    prospective_epoch_id: EpochId,
) -> anyhow::Result<()> {
    vote_new_parameters(
        contract,
        prospective_epoch_id.0,
        &new_threshold_parameters,
        remaining_accounts,
        &[],
    )
    .await?;
    conclude_resharing(contract, remaining_accounts, prospective_epoch_id).await?;
    Ok(())
}