// This file is part of the Polymesh distribution (https://github.com/PolymeshAssociation/Polymesh).
// Copyright (c) 2020 Polymesh Association

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.

// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use frame_benchmarking::benchmarks;

use pallet_compliance_manager::Module as ComplianceManager;
use polymesh_common_utilities::benchs::{user, AccountIdOf, User};
use polymesh_common_utilities::TestUtilsFn;
use polymesh_primitives::{Fund, FundDescription, PortfolioId, PortfolioNumber};

use super::*;
use crate::benchmarking::{currency, did_whts, set_ca_targets, setup_ca, SEED};
use crate::{CAKind, CorporateActions};

const MAX_TARGETS: u32 = 1000;
const MAX_DID_WHT_IDS: u32 = 1000;

fn portfolio<T: Config>(
    owner: &User<T>,
    pnum: PortfolioNumber,
    asset_id: AssetId,
    amount: Balance,
) {
    let did = owner.did();
    let origin: T::RuntimeOrigin = owner.origin().into();
    <Portfolio<T>>::create_portfolio(origin.clone(), "portfolio".into()).unwrap();
    <Portfolio<T>>::move_portfolio_funds(
        origin,
        PortfolioId::default_portfolio(did),
        PortfolioId::user_portfolio(did, pnum),
        vec![Fund {
            description: FundDescription::Fungible { asset_id, amount },
            memo: None,
        }],
    )
    .unwrap();
}

fn dist<T: Config + TestUtilsFn<AccountIdOf<T>>>(target_ids: u32) -> (User<T>, CAId, AssetId) {
    let (owner, ca_id) = setup_ca::<T>(CAKind::UnpredictableBenefit);

    let currency = currency::<T>(&owner);
    let amount = 1000u32.into();
    let pnum = 1u64.into();
    portfolio::<T>(&owner, pnum, currency, amount);

    <Module<T>>::distribute(
        owner.origin().into(),
        ca_id,
        Some(pnum),
        currency,
        2u32.into(),
        amount,
        3000,
        Some(4000),
    )
    .unwrap();

    set_ca_targets::<T>(ca_id, target_ids);

    (owner, ca_id, currency)
}

fn prepare_transfer<T: Config + pallet_compliance_manager::Config + TestUtilsFn<AccountIdOf<T>>>(
    target_ids: u32,
    did_whts_num: u32,
) -> (User<T>, User<T>, CAId) {
    let (owner, ca_id, currency) = dist::<T>(target_ids);

    CorporateActions::mutate(ca_id.asset_id, ca_id.local_id, |ca| {
        let mut whts = did_whts::<T>(did_whts_num);
        whts.sort_by_key(|(did, _)| *did);
        ca.as_mut().unwrap().withholding_tax = whts;
    });

    <pallet_timestamp::Now<T>>::set(3000u32.into());

    let holder = user::<T>("holder", SEED);
    <ComplianceManager<T>>::add_compliance_requirement(
        owner.origin().into(),
        currency,
        vec![],
        vec![],
    )
    .unwrap();

    (owner, holder, ca_id)
}

benchmarks! {
    where_clause { where
        T: pallet_compliance_manager::Config,
        T: TestUtilsFn<AccountIdOf<T>>,
    }

    distribute {
        let (owner, ca_id) = setup_ca::<T>(CAKind::UnpredictableBenefit);
        let currency = currency::<T>(&owner);
        let amount = 1000u32.into();
        let per_share = 2u32.into();
        let pnum = 1u64.into();
        portfolio::<T>(&owner, pnum, currency, amount);
    }: _(owner.origin(), ca_id, Some(pnum), currency, per_share, amount, 3000, Some(4000))
    verify {
        assert!(Distributions::get(ca_id).is_some(), "distribution not created");
    }

    claim {
        let t in 0..MAX_TARGETS;
        let w in 0..MAX_DID_WHT_IDS;

        let (_, holder, ca_id) = prepare_transfer::<T>(t, w);
    }: _(holder.origin(), ca_id)
    verify {
        assert!(HolderPaid::get((ca_id, holder.did())), "not paid");
    }

    push_benefit {
        let t in 0..MAX_TARGETS;
        let w in 0..MAX_DID_WHT_IDS;

        let (owner, holder, ca_id) = prepare_transfer::<T>(t, w);
    }: _(owner.origin(), ca_id, holder.did())
    verify {
        assert!(HolderPaid::get((ca_id, holder.did())), "not paid");
    }

    reclaim {
        let (owner, ca_id, currency) = dist::<T>(0);

        <pallet_timestamp::Now<T>>::set(5000u32.into());
    }: _(owner.origin(), ca_id)
    verify {
        assert!(Distributions::get(ca_id).unwrap().reclaimed, "not reclaimed");
    }

    remove_distribution {
        let (owner, ca_id, currency) = dist::<T>(0);
    }: _(owner.origin(), ca_id)
    verify {
        assert!(Distributions::get(ca_id).is_none(), "not removed");
    }
}
