pezframe_benchmarking::define_benchmarks!(
    [pezframe_system, SystemBench::<Runtime>]
    [pezpallet_assets, Assets]
    [pezpallet_balances, Balances]
    [pezpallet_session, SessionBench::<Runtime>]
    [pezpallet_timestamp, Timestamp]
    [pezpallet_message_queue, MessageQueue]
    [pezpallet_sudo, Sudo]
    [pezpallet_collator_selection, CollatorSelection]
    [pezcumulus_pezpallet_xcmp_queue, XcmpQueue]
    [pezcumulus_pezpallet_teyrchain_system, TeyrChainSystem]
    [pezpallet_proxy, Proxy]
    [pezpallet_utility, Utility]
    [pezpallet_multisig, Multisig]
    [pezpallet_xcm, PezpalletXcmExtrinsicsBenchmark::<Runtime>]
    [pezpallet_asset_manager, AssetManager]
    [pezpallet_referenda, Referenda]
    [pezpallet_scheduler, Scheduler]
    [pezpallet_preimage, Preimage]
    [pezpallet_treasury, Treasury]
    [pezpallet_conviction_voting, ConvictionVoting]
    [pezpallet_whitelist, Whitelist]
    [pezpallet_xcm_transactor, XcmTransactor]
    [pezpallet_xcm_weight_trader, XcmWeightTrader]
);

use pezcumulus_primitives_core::{ChannelStatus, GetChannelInfo};
use pezframe_support::traits::{
    tokens::{Pay, PaymentStatus},
    Get,
};
use pezsp_std::marker::PhantomData;

use crate::TeyrChainSystem;

/// Trait for setting up any prerequisites for successful execution of benchmarks.
pub trait EnsureSuccessful {
    fn ensure_successful();
}

/// Implementation of the [`EnsureSuccessful`] trait which opens an HRMP channel between
/// the Collectives and a teyrchain with a given ID.
pub struct OpenHrmpChannel<I>(PhantomData<I>);
impl<I: Get<u32>> EnsureSuccessful for OpenHrmpChannel<I> {
    fn ensure_successful() {
        if let ChannelStatus::Closed = TeyrChainSystem::get_channel_status(I::get().into()) {
            TeyrChainSystem::open_outbound_hrmp_channel_for_benchmarks_or_tests(I::get().into())
        }
    }
}

/// Type that wraps a type implementing the [`Pay`] trait to decorate its
/// [`Pay::ensure_successful`] function with a provided implementation of the
/// [`EnsureSuccessful`] trait.
pub struct PayWithEnsure<O, E>(PhantomData<(O, E)>);
impl<O, E> Pay for PayWithEnsure<O, E>
where
    O: Pay,
    E: EnsureSuccessful,
{
    type AssetKind = O::AssetKind;
    type Balance = O::Balance;
    type Beneficiary = O::Beneficiary;
    type Error = O::Error;
    type Id = O::Id;

    fn pay(
        who: &Self::Beneficiary,
        asset_kind: Self::AssetKind,
        amount: Self::Balance,
    ) -> Result<Self::Id, Self::Error> {
        O::pay(who, asset_kind, amount)
    }

    fn check_payment(id: Self::Id) -> PaymentStatus {
        O::check_payment(id)
    }

    fn ensure_successful(
        who: &Self::Beneficiary,
        asset_kind: Self::AssetKind,
        amount: Self::Balance,
    ) {
        E::ensure_successful();
        O::ensure_successful(who, asset_kind, amount)
    }

    fn ensure_concluded(id: Self::Id) {
        O::ensure_concluded(id)
    }
}
