use solana_sdk::{
    clock::SECONDS_PER_DAY,
    instruction::Instruction,
    message::Message,
    pubkey::Pubkey,
    stake::{
        self,
        instruction::{self as stake_instruction, LockupArgs},
        state::{Authorized, Lockup, StakeAuthorize},
    },
};

const DAYS_PER_YEAR: f64 = 365.25;
const SECONDS_PER_YEAR: i64 = (SECONDS_PER_DAY as f64 * DAYS_PER_YEAR) as i64;

pub(crate) fn derive_stake_account_address(base_pubkey: &Pubkey, i: usize) -> Pubkey {
    Pubkey::create_with_seed(base_pubkey, &i.to_string(), &stake::program::id()).unwrap()
}

// Return derived addresses
pub(crate) fn derive_stake_account_addresses(
    base_pubkey: &Pubkey,
    num_accounts: usize,
) -> Vec<Pubkey> {
    (0..num_accounts)
        .map(|i| derive_stake_account_address(base_pubkey, i))
        .collect()
}

pub(crate) fn new_stake_account(
    fee_payer_pubkey: &Pubkey,
    funding_pubkey: &Pubkey,
    base_pubkey: &Pubkey,
    lamports: u64,
    stake_authority_pubkey: &Pubkey,
    withdraw_authority_pubkey: &Pubkey,
    custodian_pubkey: &Pubkey,
    index: usize,
) -> Message {
    let stake_account_address = derive_stake_account_address(base_pubkey, index);
    let authorized = Authorized {
        staker: *stake_authority_pubkey,
        withdrawer: *withdraw_authority_pubkey,
    };
    let lockup = Lockup {
        custodian: *custodian_pubkey,
        ..Lockup::default()
    };
    let instructions = stake_instruction::create_account_with_seed(
        funding_pubkey,
        &stake_account_address,
        base_pubkey,
        &index.to_string(),
        &authorized,
        &lockup,
        lamports,
    );
    Message::new(&instructions, Some(fee_payer_pubkey))
}

fn authorize_stake_accounts_instructions(
    stake_account_address: &Pubkey,
    stake_authority_pubkey: &Pubkey,
    withdraw_authority_pubkey: &Pubkey,
    new_stake_authority_pubkey: &Pubkey,
    new_withdraw_authority_pubkey: &Pubkey,
) -> Vec<Instruction> {
    let instruction0 = stake_instruction::authorize(
        stake_account_address,
        stake_authority_pubkey,
        new_stake_authority_pubkey,
        StakeAuthorize::Staker,
        None,
    );
    let instruction1 = stake_instruction::authorize(
        stake_account_address,
        withdraw_authority_pubkey,
        new_withdraw_authority_pubkey,
        StakeAuthorize::Withdrawer,
        None,
    );
    vec![instruction0, instruction1]
}

fn rebase_stake_account(
    stake_account_address: &Pubkey,
    new_base_pubkey: &Pubkey,
    i: usize,
    fee_payer_pubkey: &Pubkey,
    stake_authority_pubkey: &Pubkey,
    lamports: u64,
) -> Option<Message> {
    if lamports == 0 {
        return None;
    }
    let new_stake_account_address = derive_stake_account_address(new_base_pubkey, i);
    let instructions = stake_instruction::split_with_seed(
        stake_account_address,
        stake_authority_pubkey,
        lamports,
        &new_stake_account_address,
        new_base_pubkey,
        &i.to_string(),
    );
    let message = Message::new(&instructions, Some(fee_payer_pubkey));
    Some(message)
}

fn move_stake_account(
    stake_account_address: &Pubkey,
    new_base_pubkey: &Pubkey,
    i: usize,
    fee_payer_pubkey: &Pubkey,
    stake_authority_pubkey: &Pubkey,
    withdraw_authority_pubkey: &Pubkey,
    new_stake_authority_pubkey: &Pubkey,
    new_withdraw_authority_pubkey: &Pubkey,
    lamports: u64,
) -> Option<Message> {
    if lamports == 0 {
        return None;
    }
    let new_stake_account_address = derive_stake_account_address(new_base_pubkey, i);
    let mut instructions = stake_instruction::split_with_seed(
        stake_account_address,
        stake_authority_pubkey,
        lamports,
        &new_stake_account_address,
        new_base_pubkey,
        &i.to_string(),
    );

    let authorize_instructions = authorize_stake_accounts_instructions(
        &new_stake_account_address,
        stake_authority_pubkey,
        withdraw_authority_pubkey,
        new_stake_authority_pubkey,
        new_withdraw_authority_pubkey,
    );

    instructions.extend(authorize_instructions.into_iter());
    let message = Message::new(&instructions, Some(fee_payer_pubkey));
    Some(message)
}

pub(crate) fn authorize_stake_accounts(
    fee_payer_pubkey: &Pubkey,
    base_pubkey: &Pubkey,
    stake_authority_pubkey: &Pubkey,
    withdraw_authority_pubkey: &Pubkey,
    new_stake_authority_pubkey: &Pubkey,
    new_withdraw_authority_pubkey: &Pubkey,
    num_accounts: usize,
) -> Vec<Message> {
    let stake_account_addresses = derive_stake_account_addresses(base_pubkey, num_accounts);
    stake_account_addresses
        .iter()
        .map(|stake_account_address| {
            let instructions = authorize_stake_accounts_instructions(
                stake_account_address,
                stake_authority_pubkey,
                withdraw_authority_pubkey,
                new_stake_authority_pubkey,
                new_withdraw_authority_pubkey,
            );
            Message::new(&instructions, Some(fee_payer_pubkey))
        })
        .collect::<Vec<_>>()
}

fn extend_lockup(lockup: &LockupArgs, years: f64) -> LockupArgs {
    let offset = (SECONDS_PER_YEAR as f64 * years) as i64;
    let unix_timestamp = lockup.unix_timestamp.map(|x| x + offset);
    let epoch = lockup.epoch.map(|_| todo!());
    LockupArgs {
        unix_timestamp,
        epoch,
        custodian: lockup.custodian,
    }
}

fn apply_lockup_changes(lockup: &LockupArgs, existing_lockup: &Lockup) -> LockupArgs {
    let custodian = match lockup.custodian {
        Some(x) if x == existing_lockup.custodian => None,
        x => x,
    };
    let epoch = match lockup.epoch {
        Some(x) if x == existing_lockup.epoch => None,
        x => x,
    };
    let unix_timestamp = match lockup.unix_timestamp {
        Some(x) if x == existing_lockup.unix_timestamp => None,
        x => x,
    };
    LockupArgs {
        unix_timestamp,
        epoch,
        custodian,
    }
}

pub(crate) fn lockup_stake_accounts(
    fee_payer_pubkey: &Pubkey,
    custodian_pubkey: &Pubkey,
    lockup: &LockupArgs,
    existing_lockups: &[(Pubkey, Lockup)],
    unlock_years: Option<f64>,
) -> Vec<Message> {
    let default_lockup = LockupArgs::default();
    existing_lockups
        .iter()
        .enumerate()
        .filter_map(|(index, (address, existing_lockup))| {
            let lockup = if let Some(unlock_years) = unlock_years {
                let unlocks = existing_lockups.len() - 1;
                let years = (unlock_years / unlocks as f64) * index as f64;
                extend_lockup(lockup, years)
            } else {
                *lockup
            };
            let lockup = apply_lockup_changes(&lockup, existing_lockup);
            if lockup == default_lockup {
                return None;
            }
            let instruction = stake_instruction::set_lockup(address, &lockup, custodian_pubkey);
            let message = Message::new(&[instruction], Some(fee_payer_pubkey));
            Some(message)
        })
        .collect()
}

pub(crate) fn rebase_stake_accounts(
    fee_payer_pubkey: &Pubkey,
    new_base_pubkey: &Pubkey,
    stake_authority_pubkey: &Pubkey,
    balances: &[(Pubkey, u64)],
) -> Vec<Message> {
    balances
        .iter()
        .enumerate()
        .filter_map(|(i, (stake_account_address, lamports))| {
            rebase_stake_account(
                stake_account_address,
                new_base_pubkey,
                i,
                fee_payer_pubkey,
                stake_authority_pubkey,
                *lamports,
            )
        })
        .collect()
}

pub(crate) fn move_stake_accounts(
    fee_payer_pubkey: &Pubkey,
    new_base_pubkey: &Pubkey,
    stake_authority_pubkey: &Pubkey,
    withdraw_authority_pubkey: &Pubkey,
    new_stake_authority_pubkey: &Pubkey,
    new_withdraw_authority_pubkey: &Pubkey,
    balances: &[(Pubkey, u64)],
) -> Vec<Message> {
    balances
        .iter()
        .enumerate()
        .filter_map(|(i, (stake_account_address, lamports))| {
            move_stake_account(
                stake_account_address,
                new_base_pubkey,
                i,
                fee_payer_pubkey,
                stake_authority_pubkey,
                withdraw_authority_pubkey,
                new_stake_authority_pubkey,
                new_withdraw_authority_pubkey,
                *lamports,
            )
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use {
        super::*,
        solana_runtime::{bank::Bank, bank_client::BankClient},
        solana_sdk::{
            account::{AccountSharedData, ReadableAccount},
            client::SyncClient,
            genesis_config::create_genesis_config,
            signature::{Keypair, Signer},
            stake::state::StakeState,
        },
        solana_stake_program::stake_state,
    };

    fn create_bank(lamports: u64) -> (Bank, Keypair, u64, u64) {
        let (genesis_config, mint_keypair) = create_genesis_config(lamports);
        let bank = Bank::new_for_tests(&genesis_config);
        let stake_rent = bank.get_minimum_balance_for_rent_exemption(StakeState::size_of());
        let system_rent = bank.get_minimum_balance_for_rent_exemption(0);
        (bank, mint_keypair, stake_rent, system_rent)
    }

    fn create_account<C: SyncClient>(
        client: &C,
        funding_keypair: &Keypair,
        lamports: u64,
    ) -> Keypair {
        let fee_payer_keypair = Keypair::new();
        client
            .transfer_and_confirm(lamports, funding_keypair, &fee_payer_keypair.pubkey())
            .unwrap();
        fee_payer_keypair
    }

    fn get_account_at<C: SyncClient>(
        client: &C,
        base_pubkey: &Pubkey,
        i: usize,
    ) -> AccountSharedData {
        let account_address = derive_stake_account_address(base_pubkey, i);
        AccountSharedData::from(client.get_account(&account_address).unwrap().unwrap())
    }

    fn get_balances<C: SyncClient>(
        client: &C,
        base_pubkey: &Pubkey,
        num_accounts: usize,
    ) -> Vec<(Pubkey, u64)> {
        (0..num_accounts)
            .map(|i| {
                let address = derive_stake_account_address(base_pubkey, i);
                (address, client.get_balance(&address).unwrap())
            })
            .collect()
    }

    fn get_lockups<C: SyncClient>(
        client: &C,
        base_pubkey: &Pubkey,
        num_accounts: usize,
    ) -> Vec<(Pubkey, Lockup)> {
        (0..num_accounts)
            .map(|i| {
                let address = derive_stake_account_address(base_pubkey, i);
                let account =
                    AccountSharedData::from(client.get_account(&address).unwrap().unwrap());
                (address, stake_state::lockup_from(&account).unwrap())
            })
            .collect()
    }

    #[test]
    fn test_new_derived_stake_account() {
        let (bank, funding_keypair, stake_rent, system_rent) = create_bank(10_000_000);
        let funding_pubkey = funding_keypair.pubkey();
        let bank_client = BankClient::new(bank);
        let fee_payer_keypair = create_account(&bank_client, &funding_keypair, system_rent);
        let fee_payer_pubkey = fee_payer_keypair.pubkey();

        let base_keypair = Keypair::new();
        let base_pubkey = base_keypair.pubkey();
        let lamports = stake_rent + 1;
        let stake_authority_pubkey = solana_sdk::pubkey::new_rand();
        let withdraw_authority_pubkey = solana_sdk::pubkey::new_rand();

        let message = new_stake_account(
            &fee_payer_pubkey,
            &funding_pubkey,
            &base_pubkey,
            lamports,
            &stake_authority_pubkey,
            &withdraw_authority_pubkey,
            &Pubkey::default(),
            0,
        );

        let signers = [&funding_keypair, &fee_payer_keypair, &base_keypair];
        bank_client
            .send_and_confirm_message(&signers, message)
            .unwrap();

        let account = get_account_at(&bank_client, &base_pubkey, 0);
        assert_eq!(account.lamports(), lamports);
        let authorized = stake_state::authorized_from(&account).unwrap();
        assert_eq!(authorized.staker, stake_authority_pubkey);
        assert_eq!(authorized.withdrawer, withdraw_authority_pubkey);
    }

    #[test]
    fn test_authorize_stake_accounts() {
        let (bank, funding_keypair, stake_rent, system_rent) = create_bank(10_000_000);
        let funding_pubkey = funding_keypair.pubkey();
        let bank_client = BankClient::new(bank);
        let fee_payer_keypair = create_account(&bank_client, &funding_keypair, system_rent);
        let fee_payer_pubkey = fee_payer_keypair.pubkey();

        let base_keypair = Keypair::new();
        let base_pubkey = base_keypair.pubkey();
        let lamports = stake_rent + 1;

        let stake_authority_keypair = Keypair::new();
        let stake_authority_pubkey = stake_authority_keypair.pubkey();
        let withdraw_authority_keypair = Keypair::new();
        let withdraw_authority_pubkey = withdraw_authority_keypair.pubkey();

        let message = new_stake_account(
            &fee_payer_pubkey,
            &funding_pubkey,
            &base_pubkey,
            lamports,
            &stake_authority_pubkey,
            &withdraw_authority_pubkey,
            &Pubkey::default(),
            0,
        );

        let signers = [&funding_keypair, &fee_payer_keypair, &base_keypair];
        bank_client
            .send_and_confirm_message(&signers, message)
            .unwrap();

        let new_stake_authority_pubkey = solana_sdk::pubkey::new_rand();
        let new_withdraw_authority_pubkey = solana_sdk::pubkey::new_rand();
        let messages = authorize_stake_accounts(
            &fee_payer_pubkey,
            &base_pubkey,
            &stake_authority_pubkey,
            &withdraw_authority_pubkey,
            &new_stake_authority_pubkey,
            &new_withdraw_authority_pubkey,
            1,
        );

        let signers = [
            &fee_payer_keypair,
            &stake_authority_keypair,
            &withdraw_authority_keypair,
        ];
        for message in messages {
            bank_client
                .send_and_confirm_message(&signers, message)
                .unwrap();
        }

        let account = get_account_at(&bank_client, &base_pubkey, 0);
        let authorized = stake_state::authorized_from(&account).unwrap();
        assert_eq!(authorized.staker, new_stake_authority_pubkey);
        assert_eq!(authorized.withdrawer, new_withdraw_authority_pubkey);
    }

    #[test]
    fn test_lockup_stake_accounts() {
        let (bank, funding_keypair, stake_rent, system_rent) = create_bank(10_000_000);
        let funding_pubkey = funding_keypair.pubkey();
        let bank_client = BankClient::new(bank);
        let fee_payer_keypair = create_account(&bank_client, &funding_keypair, system_rent);
        let fee_payer_pubkey = fee_payer_keypair.pubkey();

        let base_keypair = Keypair::new();
        let base_pubkey = base_keypair.pubkey();
        let lamports = stake_rent + 1;

        let custodian_keypair = Keypair::new();
        let custodian_pubkey = custodian_keypair.pubkey();

        let withdrawer_keypair = Keypair::new();
        let withdrawer_pubkey = withdrawer_keypair.pubkey();

        let message = new_stake_account(
            &fee_payer_pubkey,
            &funding_pubkey,
            &base_pubkey,
            lamports,
            &Pubkey::default(),
            &withdrawer_pubkey,
            &custodian_pubkey,
            0,
        );

        let signers = [&funding_keypair, &fee_payer_keypair, &base_keypair];
        bank_client
            .send_and_confirm_message(&signers, message)
            .unwrap();

        let lockups = get_lockups(&bank_client, &base_pubkey, 1);
        let messages = lockup_stake_accounts(
            &fee_payer_pubkey,
            &withdrawer_pubkey,
            &LockupArgs {
                unix_timestamp: Some(1),
                custodian: Some(custodian_pubkey),
                ..LockupArgs::default()
            },
            &lockups,
            None,
        );

        let signers = [&fee_payer_keypair, &withdrawer_keypair];
        for message in messages {
            bank_client
                .send_and_confirm_message(&signers, message)
                .unwrap();
        }

        let account = get_account_at(&bank_client, &base_pubkey, 0);
        let lockup = stake_state::lockup_from(&account).unwrap();
        assert_eq!(lockup.custodian, custodian_pubkey);
        assert_eq!(lockup.unix_timestamp, 1);
        assert_eq!(lockup.epoch, 0);

        // Assert no work left to do
        let lockups = get_lockups(&bank_client, &base_pubkey, 1);
        let messages = lockup_stake_accounts(
            &fee_payer_pubkey,
            &custodian_pubkey,
            &LockupArgs {
                unix_timestamp: Some(1),
                ..LockupArgs::default()
            },
            &lockups,
            None,
        );
        assert_eq!(messages.len(), 0);
    }

    #[test]
    fn test_rebase_empty_account() {
        let pubkey = Pubkey::default();
        let message = rebase_stake_account(&pubkey, &pubkey, 0, &pubkey, &pubkey, 0);
        assert_eq!(message, None);
    }

    #[test]
    fn test_move_empty_account() {
        let pubkey = Pubkey::default();
        let message = move_stake_account(
            &pubkey, &pubkey, 0, &pubkey, &pubkey, &pubkey, &pubkey, &pubkey, 0,
        );
        assert_eq!(message, None);
    }

    #[test]
    fn test_rebase_stake_accounts() {
        let (bank, funding_keypair, stake_rent, system_rent) = create_bank(10_000_000);
        let funding_pubkey = funding_keypair.pubkey();
        let bank_client = BankClient::new(bank);
        let fee_payer_keypair = create_account(&bank_client, &funding_keypair, system_rent);
        let fee_payer_pubkey = fee_payer_keypair.pubkey();

        let base_keypair = Keypair::new();
        let base_pubkey = base_keypair.pubkey();
        let lamports = stake_rent + 1;

        let stake_authority_keypair = Keypair::new();
        let stake_authority_pubkey = stake_authority_keypair.pubkey();
        let withdraw_authority_keypair = Keypair::new();
        let withdraw_authority_pubkey = withdraw_authority_keypair.pubkey();

        let num_accounts = 1;
        let message = new_stake_account(
            &fee_payer_pubkey,
            &funding_pubkey,
            &base_pubkey,
            lamports,
            &stake_authority_pubkey,
            &withdraw_authority_pubkey,
            &Pubkey::default(),
            0,
        );

        let signers = [&funding_keypair, &fee_payer_keypair, &base_keypair];
        bank_client
            .send_and_confirm_message(&signers, message)
            .unwrap();

        let new_base_keypair = Keypair::new();
        let new_base_pubkey = new_base_keypair.pubkey();
        let balances = get_balances(&bank_client, &base_pubkey, num_accounts);
        let messages = rebase_stake_accounts(
            &fee_payer_pubkey,
            &new_base_pubkey,
            &stake_authority_pubkey,
            &balances,
        );
        assert_eq!(messages.len(), num_accounts);

        let signers = [
            &fee_payer_keypair,
            &new_base_keypair,
            &stake_authority_keypair,
        ];
        for message in messages {
            bank_client
                .send_and_confirm_message(&signers, message)
                .unwrap();
        }

        // Ensure the new accounts are duplicates of the previous ones.
        let account = get_account_at(&bank_client, &new_base_pubkey, 0);
        let authorized = stake_state::authorized_from(&account).unwrap();
        assert_eq!(authorized.staker, stake_authority_pubkey);
        assert_eq!(authorized.withdrawer, withdraw_authority_pubkey);
    }

    #[test]
    fn test_move_stake_accounts() {
        let (bank, funding_keypair, stake_rent, system_rent) = create_bank(10_000_000);
        let funding_pubkey = funding_keypair.pubkey();
        let bank_client = BankClient::new(bank);
        let fee_payer_keypair = create_account(&bank_client, &funding_keypair, system_rent);
        let fee_payer_pubkey = fee_payer_keypair.pubkey();

        let base_keypair = Keypair::new();
        let base_pubkey = base_keypair.pubkey();
        let lamports = stake_rent + 1;

        let stake_authority_keypair = Keypair::new();
        let stake_authority_pubkey = stake_authority_keypair.pubkey();
        let withdraw_authority_keypair = Keypair::new();
        let withdraw_authority_pubkey = withdraw_authority_keypair.pubkey();

        let num_accounts = 1;
        let message = new_stake_account(
            &fee_payer_pubkey,
            &funding_pubkey,
            &base_pubkey,
            lamports,
            &stake_authority_pubkey,
            &withdraw_authority_pubkey,
            &Pubkey::default(),
            0,
        );

        let signers = [&funding_keypair, &fee_payer_keypair, &base_keypair];
        bank_client
            .send_and_confirm_message(&signers, message)
            .unwrap();

        let new_base_keypair = Keypair::new();
        let new_base_pubkey = new_base_keypair.pubkey();
        let new_stake_authority_pubkey = solana_sdk::pubkey::new_rand();
        let new_withdraw_authority_pubkey = solana_sdk::pubkey::new_rand();
        let balances = get_balances(&bank_client, &base_pubkey, num_accounts);
        let messages = move_stake_accounts(
            &fee_payer_pubkey,
            &new_base_pubkey,
            &stake_authority_pubkey,
            &withdraw_authority_pubkey,
            &new_stake_authority_pubkey,
            &new_withdraw_authority_pubkey,
            &balances,
        );
        assert_eq!(messages.len(), num_accounts);

        let signers = [
            &fee_payer_keypair,
            &new_base_keypair,
            &stake_authority_keypair,
            &withdraw_authority_keypair,
        ];
        for message in messages {
            bank_client
                .send_and_confirm_message(&signers, message)
                .unwrap();
        }

        // Ensure the new accounts have the new authorities.
        let account = get_account_at(&bank_client, &new_base_pubkey, 0);
        let authorized = stake_state::authorized_from(&account).unwrap();
        assert_eq!(authorized.staker, new_stake_authority_pubkey);
        assert_eq!(authorized.withdrawer, new_withdraw_authority_pubkey);
    }

    #[test]
    fn test_extend_lockup() {
        let lockup = LockupArgs {
            unix_timestamp: Some(1),
            ..LockupArgs::default()
        };
        let expected_lockup = LockupArgs {
            unix_timestamp: Some(1 + SECONDS_PER_YEAR),
            ..LockupArgs::default()
        };
        assert_eq!(extend_lockup(&lockup, 1.0), expected_lockup);
    }
}