use std::collections::BTreeMap;

use blstrs::G1Projective;
use group::Group;
use near_account_id::AccountId;
use near_mpc_bounded_collections::BoundedVec;
use near_mpc_contract_interface::{
    method_names,
    types::{
        Bls12381G1PublicKey, CKDAppPublicKey, CKDRequestArgs, DomainConfig,
        EDDSA_PAYLOAD_SIZE_LOWER_BOUND_BYTES, EDDSA_PAYLOAD_SIZE_UPPER_BOUND_BYTES, Payload,
        Protocol, SignRequestArgs,
    },
};
use near_primitives::action::Action;
use near_primitives::types::{Balance, Gas};
use rand::Rng;
use rand::RngCore;
use rand::rngs::OsRng;
use serde::Serialize;

/// Gas attached to a `sign` (or legacy `sign`) request. Matches the e2e
/// test cluster's `SIGN_GAS` and the contract's
/// `sign_call_gas_attachment_requirement_tera_gas` minimum.
const SIGN_TGAS: u64 = 15;
/// Gas attached to a `request_app_private_key` (CKD) call. Matches the
/// e2e cluster's `CKD_PV_GAS`. CKD is more expensive than `sign` because
/// `AppPublicKeyPV` does an on-chain bls12381 pairing check before
/// yielding.
const CKD_TGAS: u64 = 100;
/// Gas attached to a `make_parallel_sign_calls` invocation on the
/// parallel-sign helper contract. The helper schedules up to ~10
/// sub-calls into the MPC contract, so this needs the full block budget.
const PARALLEL_SIGN_TGAS: u64 = 300;

#[derive(Clone)]
pub struct ActionCall {
    pub receiver_id: AccountId,
    pub actions: Vec<Action>,
}

#[derive(Clone)]
pub struct ParallelSignCallArgs {
    pub parallel_sign_contract: AccountId,
    pub mpc_contract: AccountId,
    pub calls_by_domain: Vec<(DomainConfig, u64)>,
}

#[derive(Clone)]
pub struct RequestActionCallArgs {
    pub mpc_contract: AccountId,
    pub domain_config: DomainConfig,
}

#[derive(Clone)]
pub struct LegacySignActionCallArgs {
    pub mpc_contract: AccountId,
}

#[derive(Clone)]
pub enum ContractActionCall {
    ParallelSignCall(ParallelSignCallArgs),
    Sign(RequestActionCallArgs),
    LegacySign(LegacySignActionCallArgs),
    Ckd(RequestActionCallArgs),
}

pub fn make_actions(call: ContractActionCall) -> ActionCall {
    match call {
        ContractActionCall::ParallelSignCall(args) => {
            let mut ecdsa_calls_by_domain = BTreeMap::new();
            let mut robust_ecdsa_calls_by_domain = BTreeMap::new();
            let mut eddsa_calls_by_domain = BTreeMap::new();
            let mut ckd_calls_by_domain = BTreeMap::new();
            for (domain, prot_calls) in args.calls_by_domain {
                match domain.protocol {
                    Protocol::CaitSith => {
                        ecdsa_calls_by_domain.insert(domain.id.0, prot_calls);
                    }
                    Protocol::DamgardEtAl => {
                        robust_ecdsa_calls_by_domain.insert(domain.id.0, prot_calls);
                    }
                    Protocol::Frost => {
                        eddsa_calls_by_domain.insert(domain.id.0, prot_calls);
                    }
                    Protocol::ConfidentialKeyDerivation => {
                        ckd_calls_by_domain.insert(domain.id.0, prot_calls);
                    }
                }
            }
            ActionCall {
                receiver_id: args.parallel_sign_contract,
                actions: vec![make_action(
                    "make_parallel_sign_calls",
                    &serde_json::to_vec(&ParallelSignArgsV2 {
                        target_contract: args.mpc_contract,
                        ecdsa_calls_by_domain,
                        robust_ecdsa_calls_by_domain,
                        eddsa_calls_by_domain,
                        ckd_calls_by_domain,
                        seed: rand::random(),
                    })
                    .unwrap(),
                    PARALLEL_SIGN_TGAS,
                    1,
                )],
            }
        }
        ContractActionCall::Sign(args) => ActionCall {
            receiver_id: args.mpc_contract,
            actions: vec![make_action(
                method_names::SIGN,
                &serde_json::to_vec(&SignArgsV2 {
                    request: SignRequestArgs {
                        domain_id: args.domain_config.id,
                        path: "".to_string(),
                        payload: make_payload(args.domain_config.protocol),
                    },
                })
                .unwrap(),
                SIGN_TGAS,
                1,
            )],
        },
        ContractActionCall::LegacySign(args) => ActionCall {
            receiver_id: args.mpc_contract,
            actions: vec![make_action(
                method_names::SIGN,
                &serde_json::to_vec(&SignArgsV1 {
                    request: SignRequestV1 {
                        key_version: 0,
                        path: "".to_string(),
                        payload: rand::random(),
                    },
                })
                .unwrap(),
                SIGN_TGAS,
                1,
            )],
        },
        ContractActionCall::Ckd(args) => ActionCall {
            receiver_id: args.mpc_contract,
            actions: vec![make_action(
                method_names::REQUEST_APP_PRIVATE_KEY,
                &serde_json::to_vec(&CKDArgs {
                    request: CKDRequestArgs {
                        derivation_path: "".to_string(),
                        domain_id: args.domain_config.id,
                        app_public_key: CKDAppPublicKey::AppPublicKey(random_app_public_key()),
                    },
                })
                .unwrap(),
                CKD_TGAS,
                1,
            )],
        },
    }
}

#[derive(Serialize)]
struct SignArgsV1 {
    pub request: SignRequestV1,
}

#[derive(Serialize)]
struct SignRequestV1 {
    payload: [u8; 32],
    path: String,
    key_version: u32,
}

#[derive(Serialize)]
struct SignArgsV2 {
    pub request: SignRequestArgs,
}

#[derive(Serialize)]
struct CKDArgs {
    pub request: CKDRequestArgs,
}

#[derive(Serialize)]
struct ParallelSignArgsV2 {
    target_contract: AccountId,
    ecdsa_calls_by_domain: BTreeMap<u64, u64>,
    robust_ecdsa_calls_by_domain: BTreeMap<u64, u64>,
    eddsa_calls_by_domain: BTreeMap<u64, u64>,
    ckd_calls_by_domain: BTreeMap<u64, u64>,
    seed: u64,
}

fn make_payload(protocol: Protocol) -> Payload {
    match protocol {
        Protocol::CaitSith | Protocol::DamgardEtAl => {
            Payload::Ecdsa(rand::random::<[u8; 32]>().into())
        }
        Protocol::Frost => {
            let mut rng = rand::thread_rng();
            let len = rng.gen_range(
                EDDSA_PAYLOAD_SIZE_LOWER_BOUND_BYTES..=EDDSA_PAYLOAD_SIZE_UPPER_BOUND_BYTES,
            );
            let mut payload = vec![0; len];
            rng.fill_bytes(&mut payload);

            let bounded_payload: BoundedVec<
                u8,
                EDDSA_PAYLOAD_SIZE_LOWER_BOUND_BYTES,
                EDDSA_PAYLOAD_SIZE_UPPER_BOUND_BYTES,
            > = payload.try_into().unwrap();

            Payload::Eddsa(bounded_payload)
        }
        Protocol::ConfidentialKeyDerivation => {
            unreachable!(
                "make_payload should not be called with `ConfidentialKeyDerivation` protocol"
            )
        }
    }
}

fn random_app_public_key() -> Bls12381G1PublicKey {
    let point = G1Projective::random(&mut OsRng);
    (&point).into()
}

fn make_action(method: &str, args: &[u8], tgas: u64, deposit: u128) -> Action {
    Action::FunctionCall(Box::new(near_primitives::action::FunctionCallAction {
        method_name: method.to_string(),
        args: args.to_vec(),
        gas: Gas::from_teragas(tgas),
        deposit: Balance::from_yoctonear(deposit),
    }))
}