//! Calculate and collect rent from accounts.
#![cfg_attr(docsrs, feature(doc_auto_cfg))]
#![cfg_attr(feature = "frozen-abi", feature(min_specialization))]

use rialo_s_account::{AccountSharedData, ReadableAccount, WritableAccount};
use rialo_s_clock::Epoch;
use rialo_s_epoch_schedule::EpochSchedule;
use rialo_s_genesis_config::GenesisConfig;
use rialo_s_pubkey::Pubkey;
use rialo_s_rent::{Rent, RentDue};
use rialo_s_sdk_ids::incinerator;

#[cfg_attr(feature = "frozen-abi", derive(rialo_frozen_abi_macro::AbiExample))]
#[cfg_attr(
    feature = "serde",
    derive(serde_derive::Deserialize, serde_derive::Serialize)
)]
#[derive(Clone, Debug, PartialEq)]
pub struct RentCollector {
    pub epoch: Epoch,
    pub epoch_schedule: EpochSchedule,
    pub slots_per_year: f64,
    pub rent: Rent,
}

impl Default for RentCollector {
    fn default() -> Self {
        Self {
            epoch: Epoch::default(),
            epoch_schedule: EpochSchedule::default(),
            // derive default value using GenesisConfig::default()
            slots_per_year: GenesisConfig::default().slots_per_year(),
            rent: Rent::default(),
        }
    }
}

/// When rent is collected from an exempt account, rent_epoch is set to this
/// value. The idea is to have a fixed, consistent value for rent_epoch for all accounts that do not collect rent.
/// This enables us to get rid of the field completely.
pub const RENT_EXEMPT_RENT_EPOCH: Epoch = Epoch::MAX;

/// when rent is collected for this account, this is the action to apply to the account
#[derive(Debug)]
enum RentResult {
    /// this account will never have rent collected from it
    Exempt,
    /// maybe we collect rent later, but not now
    NoRentCollectionNow,
    /// collect rent
    CollectRent {
        new_rent_epoch: Epoch,
        rent_due: u64, // kelvins, could be 0
    },
}

impl RentCollector {
    pub fn new(
        epoch: Epoch,
        epoch_schedule: EpochSchedule,
        slots_per_year: f64,
        rent: Rent,
    ) -> Self {
        Self {
            epoch,
            epoch_schedule,
            slots_per_year,
            rent,
        }
    }

    pub fn clone_with_epoch(&self, epoch: Epoch) -> Self {
        Self {
            epoch,
            ..self.clone()
        }
    }

    /// true if it is easy to determine this account should consider having rent collected from it
    pub fn should_collect_rent(&self, address: &Pubkey, executable: bool) -> bool {
        !(executable // executable accounts must be rent-exempt balance
            || *address == incinerator::id())
    }

    /// given an account that 'should_collect_rent'
    /// returns (amount rent due, is_exempt_from_rent)
    pub fn get_rent_due(
        &self,
        kelvins: u64,
        data_len: usize,
        account_rent_epoch: Epoch,
    ) -> RentDue {
        if self.rent.is_exempt(kelvins, data_len) {
            RentDue::Exempt
        } else {
            let slots_elapsed: u64 = (account_rent_epoch..=self.epoch)
                .map(|epoch| {
                    self.epoch_schedule
                        .get_slots_in_epoch(epoch.saturating_add(1))
                })
                .sum();

            // avoid infinite rent in rust 1.45
            let years_elapsed = if self.slots_per_year != 0.0 {
                slots_elapsed as f64 / self.slots_per_year
            } else {
                0.0
            };

            // we know this account is not exempt
            let due = self.rent.due_amount(data_len, years_elapsed);
            RentDue::Paying(due)
        }
    }

    // Updates the account's kelvins and status, and returns the amount of rent collected, if any.
    // This is NOT thread safe at some level. If we try to collect from the same account in
    // parallel, we may collect twice.
    #[must_use = "add to Bank::collected_rent"]
    pub fn collect_from_existing_account(
        &self,
        address: &Pubkey,
        account: &mut AccountSharedData,
    ) -> CollectedInfo {
        // AccountSharedData alone is not executable - executables are stored as ExecutableAccount
        // in StoredAccount::Executable variant
        match self.calculate_rent_result(address, account, false) {
            RentResult::Exempt => {
                account.set_rent_epoch(RENT_EXEMPT_RENT_EPOCH);
                CollectedInfo::default()
            }
            RentResult::NoRentCollectionNow => CollectedInfo::default(),
            RentResult::CollectRent {
                new_rent_epoch,
                rent_due,
            } => match account.kelvins().checked_sub(rent_due) {
                None | Some(0) => {
                    let account = std::mem::take(account);
                    CollectedInfo {
                        rent_amount: account.kelvins(),
                        account_data_len_reclaimed: account.data().len() as u64,
                    }
                }
                Some(kelvins) => {
                    account.set_kelvins(kelvins);
                    account.set_rent_epoch(new_rent_epoch);
                    CollectedInfo {
                        rent_amount: rent_due,
                        account_data_len_reclaimed: 0u64,
                    }
                }
            },
        }
    }

    /// determine what should happen to collect rent from this account
    ///
    /// `is_executable` should be true if the account is an executable program account.
    /// This is passed separately because `ReadableAccount` doesn't track executable status -
    /// use `StoredAccount::is_executable()` to get this value.
    #[must_use]
    fn calculate_rent_result(
        &self,
        address: &Pubkey,
        account: &impl ReadableAccount,
        is_executable: bool,
    ) -> RentResult {
        if account.rent_epoch() == RENT_EXEMPT_RENT_EPOCH || account.rent_epoch() > self.epoch {
            // potentially rent paying account (or known and already marked exempt)
            // Maybe collect rent later, leave account alone for now.
            return RentResult::NoRentCollectionNow;
        }
        if !self.should_collect_rent(address, is_executable) {
            // easy to determine this account should not consider having rent collected from it
            return RentResult::Exempt;
        }
        match self.get_rent_due(
            account.kelvins(),
            account.data().len(),
            account.rent_epoch(),
        ) {
            // account will not have rent collected ever
            RentDue::Exempt => RentResult::Exempt,
            // potentially rent paying account
            // Maybe collect rent later, leave account alone for now.
            RentDue::Paying(0) => RentResult::NoRentCollectionNow,
            // Rent is collected for next epoch.
            RentDue::Paying(rent_due) => RentResult::CollectRent {
                new_rent_epoch: self.epoch.saturating_add(1),
                rent_due,
            },
        }
    }
}

/// Information computed during rent collection
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
pub struct CollectedInfo {
    /// Amount of rent collected from account
    pub rent_amount: u64,
    /// Size of data reclaimed from account (happens when account's kelvins go to zero)
    pub account_data_len_reclaimed: u64,
}

impl std::ops::Add for CollectedInfo {
    type Output = Self;
    fn add(self, other: Self) -> Self {
        Self {
            rent_amount: self.rent_amount.saturating_add(other.rent_amount),
            account_data_len_reclaimed: self
                .account_data_len_reclaimed
                .saturating_add(other.account_data_len_reclaimed),
        }
    }
}

impl std::ops::AddAssign for CollectedInfo {
    #![allow(clippy::arithmetic_side_effects)]
    fn add_assign(&mut self, other: Self) {
        *self = *self + other;
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use rialo_s_account::Account;
    use rialo_s_sdk_ids::sysvar;

    use super::*;

    fn default_rent_collector_clone_with_epoch(epoch: Epoch) -> RentCollector {
        RentCollector::default().clone_with_epoch(epoch)
    }

    impl RentCollector {
        #[must_use = "add to Bank::collected_rent"]
        fn collect_from_created_account(
            &self,
            address: &Pubkey,
            account: &mut AccountSharedData,
        ) -> CollectedInfo {
            // initialize rent_epoch as created at this epoch
            account.set_rent_epoch(self.epoch);
            self.collect_from_existing_account(address, account)
        }
    }

    #[test]
    fn test_calculate_rent_result() {
        let mut rent_collector = RentCollector::default();

        let mut account = AccountSharedData::default();
        assert_matches!(
            rent_collector.calculate_rent_result(&Pubkey::default(), &account, false),
            RentResult::NoRentCollectionNow
        );
        {
            let mut account_clone = account.clone();
            assert_eq!(
                rent_collector
                    .collect_from_existing_account(&Pubkey::default(), &mut account_clone),
                CollectedInfo::default()
            );
            assert_eq!(account_clone, account);
        }

        // Test with is_executable=true to verify executables are rent-exempt
        // Note: calculate_rent_result takes is_executable as a parameter because
        // AccountSharedData doesn't track executable status - it's stored externally
        // in StoredAccount::Executable variant. So collect_from_existing_account
        // always treats accounts as non-executable.
        assert_matches!(
            rent_collector.calculate_rent_result(&Pubkey::default(), &account, true),
            RentResult::Exempt
        );
        {
            // collect_from_existing_account doesn't know the account is executable,
            // so it returns NoRentCollectionNow (same as the non-executable case above)
            let mut account_clone = account.clone();
            assert_eq!(
                rent_collector
                    .collect_from_existing_account(&Pubkey::default(), &mut account_clone),
                CollectedInfo::default()
            );
            assert_eq!(account_clone, account); // account unchanged
        }

        // Incinerator accounts are always rent-exempt regardless of executable status
        assert_matches!(
            rent_collector.calculate_rent_result(&incinerator::id(), &account, false),
            RentResult::Exempt
        );
        {
            let mut account_clone = account.clone();
            let mut account_expected = account.clone();
            account_expected.set_rent_epoch(RENT_EXEMPT_RENT_EPOCH);
            assert_eq!(
                rent_collector
                    .collect_from_existing_account(&incinerator::id(), &mut account_clone),
                CollectedInfo::default()
            );
            assert_eq!(account_clone, account_expected);
        }

        // try a few combinations of rent collector rent epoch and collecting rent
        for (rent_epoch, rent_due_expected) in [(2, 2), (3, 5)] {
            rent_collector.epoch = rent_epoch;
            account.set_kelvins(10);
            account.set_rent_epoch(1);
            let new_rent_epoch_expected = rent_collector.epoch + 1;
            assert!(
                matches!(
                    rent_collector.calculate_rent_result(&Pubkey::default(), &account, false),
                    RentResult::CollectRent{ new_rent_epoch, rent_due} if new_rent_epoch == new_rent_epoch_expected && rent_due == rent_due_expected,
                ),
                "{:?}",
                rent_collector.calculate_rent_result(&Pubkey::default(), &account, false)
            );

            {
                let mut account_clone = account.clone();
                assert_eq!(
                    rent_collector
                        .collect_from_existing_account(&Pubkey::default(), &mut account_clone),
                    CollectedInfo {
                        rent_amount: rent_due_expected,
                        account_data_len_reclaimed: 0
                    }
                );
                let mut account_expected = account.clone();
                account_expected.set_kelvins(account.kelvins() - rent_due_expected);
                account_expected.set_rent_epoch(new_rent_epoch_expected);
                assert_eq!(account_clone, account_expected);
            }
        }

        // enough kelvins to make us exempt
        account.set_kelvins(1_000_000);
        let result = rent_collector.calculate_rent_result(&Pubkey::default(), &account, false);
        assert!(matches!(result, RentResult::Exempt), "{result:?}",);
        {
            let mut account_clone = account.clone();
            let mut account_expected = account.clone();
            account_expected.set_rent_epoch(RENT_EXEMPT_RENT_EPOCH);
            assert_eq!(
                rent_collector
                    .collect_from_existing_account(&Pubkey::default(), &mut account_clone),
                CollectedInfo::default()
            );
            assert_eq!(account_clone, account_expected);
        }

        // enough kelvins to make us exempt
        // but, our rent_epoch is set in the future, so we can't know if we are exempt yet or not.
        // We don't calculate rent amount vs data if the rent_epoch is already in the future.
        account.set_rent_epoch(1_000_000);
        assert_matches!(
            rent_collector.calculate_rent_result(&Pubkey::default(), &account, false),
            RentResult::NoRentCollectionNow
        );
        {
            let mut account_clone = account.clone();
            assert_eq!(
                rent_collector
                    .collect_from_existing_account(&Pubkey::default(), &mut account_clone),
                CollectedInfo::default()
            );
            assert_eq!(account_clone, account);
        }
    }

    #[test]
    fn test_collect_from_account_created_and_existing() {
        let old_kelvins = 1000;
        let old_epoch = 1;
        let new_epoch = 2;

        let (mut created_account, mut existing_account) = {
            let account = AccountSharedData::from(Account {
                kelvins: old_kelvins,
                rent_epoch: old_epoch,
                ..Account::default()
            });

            (account.clone(), account)
        };

        let rent_collector = default_rent_collector_clone_with_epoch(new_epoch);

        // collect rent on a newly-created account
        let collected = rent_collector
            .collect_from_created_account(&rialo_s_pubkey::new_rand(), &mut created_account);
        assert!(created_account.kelvins() < old_kelvins);
        assert_eq!(
            created_account.kelvins() + collected.rent_amount,
            old_kelvins
        );
        assert_ne!(created_account.rent_epoch(), old_epoch);
        assert_eq!(collected.account_data_len_reclaimed, 0);

        // collect rent on a already-existing account
        let collected = rent_collector
            .collect_from_existing_account(&rialo_s_pubkey::new_rand(), &mut existing_account);
        assert!(existing_account.kelvins() < old_kelvins);
        assert_eq!(
            existing_account.kelvins() + collected.rent_amount,
            old_kelvins
        );
        assert_ne!(existing_account.rent_epoch(), old_epoch);
        assert_eq!(collected.account_data_len_reclaimed, 0);

        // newly created account should be collected for less rent; thus more remaining balance
        assert!(created_account.kelvins() > existing_account.kelvins());
        assert_eq!(created_account.rent_epoch(), existing_account.rent_epoch());
    }

    #[test]
    fn test_rent_exempt_temporal_escape() {
        for pass in 0..2 {
            let mut account = AccountSharedData::default();
            let epoch = 3;
            let huge_kelvins = 123_456_789_012;
            let tiny_kelvins = 789_012;
            let pubkey = rialo_s_pubkey::new_rand();

            assert_eq!(account.rent_epoch(), 0);

            // create a tested rent collector
            let rent_collector = default_rent_collector_clone_with_epoch(epoch);

            if pass == 0 {
                account.set_kelvins(huge_kelvins);
                // first mark account as being collected while being rent-exempt
                let collected = rent_collector.collect_from_existing_account(&pubkey, &mut account);
                assert_eq!(account.kelvins(), huge_kelvins);
                assert_eq!(collected, CollectedInfo::default());
                continue;
            }

            // decrease the balance not to be rent-exempt
            // In a real validator, it is not legal to reduce an account's kelvins such that the account becomes rent paying.
            // So, pass == 0 above tests the case of rent that is exempt. pass == 1 tests the case where we are rent paying.
            account.set_kelvins(tiny_kelvins);

            // ... and trigger another rent collection on the same epoch and check that rent is working
            let collected = rent_collector.collect_from_existing_account(&pubkey, &mut account);
            assert_eq!(account.kelvins(), tiny_kelvins - collected.rent_amount);
            assert_ne!(collected, CollectedInfo::default());
        }
    }

    #[test]
    fn test_rent_exempt_sysvar() {
        let tiny_kelvins = 1;
        let mut account = AccountSharedData::default();
        account.set_owner(sysvar::id());
        account.set_kelvins(tiny_kelvins);

        let pubkey = rialo_s_pubkey::new_rand();

        assert_eq!(account.rent_epoch(), 0);

        let epoch = 3;
        let rent_collector = default_rent_collector_clone_with_epoch(epoch);

        let collected = rent_collector.collect_from_existing_account(&pubkey, &mut account);
        assert_eq!(account.kelvins(), 0);
        assert_eq!(collected.rent_amount, 1);
    }

    /// Ensure that when an account is "rent collected" away, its data len is returned.
    #[test]
    fn test_collect_cleans_up_account() {
        rialo_s_logger::setup();
        let account_kelvins = 1; // must be *below* rent amount
        let account_data_len = 567;
        let account_rent_epoch = 11;
        let mut account = AccountSharedData::from(Account {
            kelvins: account_kelvins, // <-- must be below rent-exempt amount
            data: vec![u8::default(); account_data_len],
            rent_epoch: account_rent_epoch,
            ..Account::default()
        });
        let rent_collector = default_rent_collector_clone_with_epoch(account_rent_epoch + 1);

        let collected =
            rent_collector.collect_from_existing_account(&Pubkey::new_unique(), &mut account);

        assert_eq!(collected.rent_amount, account_kelvins);
        assert_eq!(
            collected.account_data_len_reclaimed,
            account_data_len as u64
        );
        assert_eq!(account, AccountSharedData::default());
    }
}
