use crate::accounts_zerocopy::*;
use crate::error::MangoError;
use crate::logs::{MarginTradeLog, TokenBalanceLog};
use crate::state::{
    compute_health, new_fixed_order_account_retriever, AccountRetriever, Bank, Group, HealthType,
    MangoAccount,
};
use crate::{group_seeds, Mango};
use anchor_lang::prelude::*;
use anchor_spl::token::{self, Token, TokenAccount};
use fixed::types::I80F48;
use solana_program::instruction::Instruction;
use std::cell::Ref;
use std::collections::HashMap;

/// The flash loan instruction
///
/// In addition to these accounts, there must be a sequence of remaining_accounts:
/// 1. health_accounts: accounts needed for health checking
/// 2. per cpi
/// 2.a. target_program_id: the target program account
/// 2.b. target_accounts: the accounts to pass to the target program
///
/// Every vault address listed in 3. must also have the matching bank and oracle appear in 1.
///
/// Every vault that is to be withdrawn from must appear in the `withdraws` instruction argument.
/// The corresponding bank may be used as an authority for vault withdrawals.
#[derive(Accounts)]
pub struct FlashLoan<'info> {
    pub group: AccountLoader<'info, Group>,

    #[account(
        mut,
        has_one = group,
        constraint = account.load()?.is_owner_or_delegate(owner.key()),
    )]
    pub account: AccountLoader<'info, MangoAccount>,

    pub owner: Signer<'info>,
    pub token_program: Program<'info, Token>,
}

#[derive(Debug)]
struct AllowedVault {
    /// index of the vault in cpi_ais
    vault_cpi_ai_index: usize,
    /// index of the bank in health_ais
    bank_health_ai_index: usize,
    /// raw index into account.tokens
    raw_token_index: usize,
    /// vault amount before cpi
    pre_amount: u64,
    /// requested withdraw amount
    withdraw_amount: u64,
    /// amount of withdraw request that is a loan
    loan_amount: I80F48,
}

#[derive(AnchorDeserialize, AnchorSerialize, Clone, Copy, Debug)]
pub struct FlashLoanWithdraw {
    /// Account index of the vault to withdraw from in the target_accounts section.
    /// Index is counted after health accounts.
    pub index: u8,
    /// Requested withdraw amount.
    pub amount: u64,
}

#[derive(AnchorDeserialize, AnchorSerialize, Debug)]
pub struct CpiData {
    pub account_start: u8,
    pub data: Vec<u8>,
}

/// - `withdraws` is a list of FlashLoanWithdraw requests.
/// - `cpi_datas` is a list of bytes per cpi to call the target_program_id with.
/// - `cpi_account_starts` is a list of index into the remaining accounts per cpi to call the target_program_id with.
pub fn flash_loan<'key, 'accounts, 'remaining, 'info>(
    ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoan<'info>>,
    withdraws: Vec<FlashLoanWithdraw>,
    cpi_datas: Vec<CpiData>,
) -> Result<()> {
    require!(!cpi_datas.is_empty(), MangoError::SomeError);
    let num_of_cpis = cpi_datas.len();
    let num_health_accounts = cpi_datas.get(0).unwrap().account_start as usize;

    let group = ctx.accounts.group.load()?;
    let mut account = ctx.accounts.account.load_mut()?;
    require!(account.is_bankrupt == 0, MangoError::IsBankrupt);

    // Go over the banks passed as health accounts and:
    // - Ensure that all banks that are passed in have activated positions.
    //   This is necessary because maybe the user wants to margin trade on a token
    //   that the account hasn't used before.
    // - Collect the addresses of all banks to potentially sign for in cpi_ais.
    // - Collect the addresses of all bank vaults.
    // Note: This depends on the particular health account ordering.
    let health_ais = &ctx.remaining_accounts[0..num_health_accounts];
    let mut allowed_banks = HashMap::<&Pubkey, Ref<Bank>>::new();
    // vault pubkey -> (bank_account_index, raw_token_index)
    let mut allowed_vaults = HashMap::<Pubkey, (usize, usize)>::new();
    for (i, ai) in health_ais.iter().enumerate() {
        match ai.load::<Bank>() {
            Ok(bank) => {
                require!(bank.group == account.group, MangoError::SomeError);
                let (_, raw_token_index, _) = account.tokens.get_mut_or_create(bank.token_index)?;
                allowed_vaults.insert(bank.vault, (i, raw_token_index));
                allowed_banks.insert(ai.key, bank);
            }
            Err(Error::AnchorError(error))
                if error.error_code_number == ErrorCode::AccountDiscriminatorMismatch as u32
                    || error.error_code_number == ErrorCode::AccountOwnedByWrongProgram as u32 =>
            {
                break;
            }
            Err(error) => return Err(error),
        };
    }

    // Check pre-cpi health
    // NOTE: This health check isn't strictly necessary. It will be, later, when
    // we want to have reduce_only or be able to move an account out of bankruptcy.
    {
        let retriever = new_fixed_order_account_retriever(health_ais, &account)?;
        let pre_cpi_health = compute_health(&account, HealthType::Init, &retriever)?;
        require!(pre_cpi_health >= 0, MangoError::HealthMustBePositive);
        msg!("pre_cpi_health {:?}", pre_cpi_health);
    }

    let all_cpi_ais = &ctx.remaining_accounts[num_health_accounts..];
    let mut all_cpi_ams = all_cpi_ais
        .iter()
        .flat_map(|item| item.to_account_metas(None))
        .collect::<Vec<_>>();
    require!(
        all_cpi_ais.len() == all_cpi_ams.len(),
        MangoError::SomeError
    );

    // Check that each group-owned token account is the vault of one of the allowed banks,
    // and track its balance.
    let mut used_vaults = all_cpi_ais
        .iter()
        .enumerate()
        .filter_map(|(i, ai)| {
            if ai.owner != &TokenAccount::owner() {
                return None;
            }

            // Skip mints and other accounts that may be owned by the spl_token program
            let maybe_token_account = Account::<TokenAccount>::try_from(ai);
            if maybe_token_account.is_err() {
                return None;
            }

            let token_account = maybe_token_account.unwrap();
            if token_account.owner != ctx.accounts.group.key() {
                return None;
            }

            // Every group-owned token account must be a vault of one of the banks.
            if let Some(&(bank_index, raw_token_index)) = allowed_vaults.get(&ai.key) {
                return Some(Ok((
                    ai.key,
                    AllowedVault {
                        vault_cpi_ai_index: i,
                        bank_health_ai_index: bank_index,
                        raw_token_index,
                        pre_amount: token_account.amount,
                        // these two are updated later
                        withdraw_amount: 0,
                        loan_amount: I80F48::ZERO,
                    },
                )));
            }

            // This is to protect users, because if their cpi program sends deposits to a vault
            // and they forgot to pass in the bank for the vault, their account would not be credited.
            Some(Err(error!(MangoError::SomeError)))
        })
        .collect::<Result<HashMap<_, _>>>()?;

    // Store the indexed value before the margin trade for logging purposes
    let mut pre_indexed_positions = Vec::new();
    for (_, info) in used_vaults.iter() {
        let position = account.tokens.get_raw(info.raw_token_index);
        pre_indexed_positions.push(position.indexed_position.to_bits());
    }

    // Find banks for used vaults in cpi_ais and collect signer seeds for them.
    // Also update withdraw_amount and loan_amount.
    let mut bank_signer_data = Vec::with_capacity(used_vaults.len());
    for (ai, am) in all_cpi_ais.iter().zip(all_cpi_ams.iter_mut()) {
        if ai.owner != &Mango::id() {
            continue;
        }
        if let Some(bank) = allowed_banks.get(ai.key) {
            if let Some(vault_info) = used_vaults.get_mut(&bank.vault) {
                let withdraw_amount = withdraws
                    .iter()
                    .find_map(|&withdraw| {
                        (withdraw.index as usize == vault_info.vault_cpi_ai_index)
                            .then(|| withdraw.amount)
                    })
                    // Even if we don't withdraw from a vault we still need to track it:
                    // Possibly the invoked program will deposit funds into it.
                    .unwrap_or(0);
                require!(
                    withdraw_amount <= vault_info.pre_amount,
                    MangoError::SomeError
                );
                vault_info.withdraw_amount = withdraw_amount;

                // if there are withdraws: figure out loan amount, mark as signer
                if withdraw_amount > 0 {
                    let token_account = account.tokens.get_mut_raw(vault_info.raw_token_index);
                    let native_position = token_account.native(&bank);
                    vault_info.loan_amount = if native_position > 0 {
                        (I80F48::from(withdraw_amount) - native_position).max(I80F48::ZERO)
                    } else {
                        I80F48::from(withdraw_amount)
                    };

                    am.is_signer = true;
                    // this is the data we'll need later to build the PDA account signer seeds
                    bank_signer_data.push((
                        bank.token_index.to_le_bytes(),
                        bank.bank_num.to_le_bytes(),
                        [bank.bump],
                    ));
                }
            }
        }
    }

    // Approve bank delegates for withdrawals
    let group_seeds = group_seeds!(group);
    let seeds = [&group_seeds[..]];
    for (_, vault_info) in used_vaults.iter() {
        if vault_info.withdraw_amount > 0 {
            let approve_ctx = CpiContext::new(
                ctx.accounts.token_program.to_account_info(),
                token::Approve {
                    to: all_cpi_ais[vault_info.vault_cpi_ai_index].clone(),
                    delegate: health_ais[vault_info.bank_health_ai_index].clone(),
                    authority: ctx.accounts.group.to_account_info(),
                },
            )
            .with_signer(&seeds);
            token::approve(approve_ctx, vault_info.withdraw_amount)?;
        }
    }

    msg!("withdraws {:#?}", withdraws);
    msg!("cpi_datas {:#?}", cpi_datas);
    msg!("allowed_vaults {:#?}", allowed_vaults);
    msg!("used_vaults {:#?}", used_vaults);

    // get rid of Ref<> to avoid limiting the cpi call
    drop(allowed_banks);
    drop(group);
    drop(account);

    // prepare signer seeds and invoke cpi
    let group_key = ctx.accounts.group.key();
    let signers = bank_signer_data
        .iter()
        .map(|(token_index, bank_num, bump)| {
            [
                group_key.as_ref(),
                b"Bank".as_ref(),
                &token_index[..],
                &bank_num[..],
                &bump[..],
            ]
        })
        .collect::<Vec<_>>();
    let signers_ref = signers.iter().map(|v| &v[..]).collect::<Vec<_>>();
    for (cpi_index, cpi_data) in cpi_datas.iter().enumerate() {
        let cpi_account_start = cpi_data.account_start as usize;
        let cpi_program_id = *ctx.remaining_accounts[cpi_account_start].key;
        require_keys_neq!(cpi_program_id, crate::id(), MangoError::SomeError);

        let all_cpi_ais_end_index = if cpi_index == num_of_cpis - 1 {
            all_cpi_ams.len()
        } else {
            cpi_datas[cpi_index + 1].account_start as usize - num_health_accounts
        };

        let all_cpi_ais_start_index = cpi_account_start - num_health_accounts + 1;

        let cpi_ais = &all_cpi_ais[all_cpi_ais_start_index..all_cpi_ais_end_index];
        let cpi_ams = &all_cpi_ams[all_cpi_ais_start_index..all_cpi_ais_end_index];
        let cpi_ix = Instruction {
            program_id: cpi_program_id,
            // todo future: optimise out these to_vecs
            data: cpi_data.data.to_vec(),
            accounts: cpi_ams.to_vec(),
        };
        solana_program::program::invoke_signed(&cpi_ix, &cpi_ais, &signers_ref)?;
    }

    // Revoke delegates for vaults
    let group = ctx.accounts.group.load()?;
    let group_seeds = group_seeds!(group);
    for (_, vault_info) in used_vaults.iter() {
        if vault_info.withdraw_amount > 0 {
            let ix = token::spl_token::instruction::revoke(
                &token::spl_token::ID,
                all_cpi_ais[vault_info.vault_cpi_ai_index].key,
                &ctx.accounts.group.key(),
                &[],
            )?;
            solana_program::program::invoke_signed(
                &ix,
                &[
                    all_cpi_ais[vault_info.vault_cpi_ai_index].clone(),
                    ctx.accounts.group.to_account_info(),
                ],
                &[group_seeds],
            )?;
        }
    }

    // Track vault changes and apply them to the user's token positions
    let mut account = ctx.accounts.account.load_mut()?;
    let inactive_tokens =
        adjust_for_post_cpi_vault_amounts(health_ais, all_cpi_ais, &used_vaults, &mut account)?;

    // Check post-cpi health
    let retriever = new_fixed_order_account_retriever(health_ais, &account)?;
    let post_cpi_health = compute_health(&account, HealthType::Init, &retriever)?;
    require!(post_cpi_health >= 0, MangoError::HealthMustBePositive);
    msg!("post_cpi_health {:?}", post_cpi_health);

    // Token balances logging
    let mut token_indexes = Vec::with_capacity(used_vaults.len());
    let mut post_indexed_positions = Vec::with_capacity(used_vaults.len());
    for (_, info) in used_vaults.iter() {
        let position = account.tokens.get_raw(info.raw_token_index);
        post_indexed_positions.push(position.indexed_position.to_bits());
        token_indexes.push(position.token_index as u16);

        let (bank, oracle_price) = retriever.bank_and_oracle(
            &ctx.accounts.group.key(),
            info.bank_health_ai_index,
            position.token_index,
        )?;

        emit!(TokenBalanceLog {
            mango_account: ctx.accounts.account.key(),
            token_index: bank.token_index as u16,
            indexed_position: position.indexed_position.to_bits(),
            deposit_index: bank.deposit_index.to_bits(),
            borrow_index: bank.borrow_index.to_bits(),
            price: oracle_price.to_bits(),
        });
    }

    emit!(MarginTradeLog {
        mango_account: ctx.accounts.account.key(),
        token_indexes,
        pre_indexed_positions,
        post_indexed_positions,
    });

    // Deactivate inactive token accounts at the end
    for raw_token_index in inactive_tokens {
        account.tokens.deactivate(raw_token_index);
    }

    Ok(())
}

fn adjust_for_post_cpi_vault_amounts(
    health_ais: &[AccountInfo],
    cpi_ais: &[AccountInfo],
    used_vaults: &HashMap<&Pubkey, AllowedVault>,
    account: &mut MangoAccount,
) -> Result<Vec<usize>> {
    let mut inactive_token_raw_indexes = Vec::with_capacity(used_vaults.len());
    for (_, info) in used_vaults.iter() {
        let vault = Account::<TokenAccount>::try_from(&cpi_ais[info.vault_cpi_ai_index]).unwrap();
        let mut bank = health_ais[info.bank_health_ai_index].load_mut::<Bank>()?;
        let position = account.tokens.get_mut_raw(info.raw_token_index);

        let loan_origination_fee = info.loan_amount * bank.loan_origination_fee_rate;
        bank.collected_fees_native += loan_origination_fee;

        let is_active = bank.change_without_fee(
            position,
            I80F48::from(vault.amount) - I80F48::from(info.pre_amount) - loan_origination_fee,
        )?;
        if !is_active {
            inactive_token_raw_indexes.push(info.raw_token_index);
        }
    }
    Ok(inactive_token_raw_indexes)
}