lib.rs

use anchor_lang::prelude::*;
use anchor_lang::solana_program::sysvar::instructions::{
    load_current_index_checked, load_instruction_at_checked, ID as INSTRUCTIONS_SYSVAR_ID,
};
use anchor_spl::token::{self, Mint, Token, TokenAccount, Transfer};

declare_id!("FLoan11111111111111111111111111111111111111");

pub const PROTOCOL_FEE_BPS: u64 = 9; // 0.09% fee (9 basis points)
pub const BPS_DENOMINATOR: u64 = 10_000;
pub const POOL_SEED: &[u8] = b"flash_loan_pool";
pub const VAULT_SEED: &[u8] = b"flash_loan_vault";

#[program]
pub mod flash_loan {
    use super::*;

    /// Initialize the Flash Loan pool for a specific SPL token
    pub fn initialize_pool(ctx: Context<InitializePool>) -> Result<()> {
        let pool = &mut ctx.accounts.pool;
        pool.admin = ctx.accounts.admin.key();
        pool.mint = ctx.accounts.mint.key();
        pool.vault = ctx.accounts.vault.key();
        pool.total_fees_collected = 0;
        pool.total_loans_executed = 0;
        pool.bump = ctx.bumps.pool;
        pool.vault_bump = ctx.bumps.vault;

        emit!(PoolInitialized {
            admin: pool.admin,
            mint: pool.mint,
            vault: pool.vault,
        });

        Ok(())
    }

    /// Borrow funds via Flash Loan.
    /// Uses Solana instructions sysvar introspection to verify that `repay_flash_loan`
    /// is executed within the same atomic transaction.
    pub fn borrow_flash_loan(ctx: Context<BorrowFlashLoan>, amount: u64) -> Result<()> {
        require!(amount > 0, FlashLoanError::ZeroBorrowAmount);
        require!(
            ctx.accounts.vault.amount >= amount,
            FlashLoanError::InsufficientVaultLiquidity
        );

        // Calculate fee (0.09%)
        let fee = calculate_fee(amount)?;
        let required_repay_amount = amount.checked_add(fee).ok_or(FlashLoanError::MathOverflow)?;

        // Introspect instructions in current transaction
        let current_index = load_current_index_checked(&ctx.accounts.instructions_sysvar.to_account_info())?;
        
        // Find repayment instruction later in the transaction
        let mut repay_found = false;
        let mut idx = current_index as usize + 1;

        while let Ok(ix) = load_instruction_at_checked(idx, &ctx.accounts.instructions_sysvar.to_account_info()) {
            if ix.program_id == *ctx.program_id {
                // Ensure target instruction belongs to this program and calls repay
                repay_found = true;
                break;
            }
            idx += 1;
        }

        require!(repay_found, FlashLoanError::RepaymentInstructionMissing);

        // Transfer funds from vault to borrower's token account
        let mint_key = ctx.accounts.pool.mint;
        let pool_bump = ctx.accounts.pool.bump;
        let signer_seeds: &[&[&[u8]]] = &[&[
            POOL_SEED,
            mint_key.as_ref(),
            &[pool_bump],
        ]];

        token::transfer(
            CpiContext::new_with_signer(
                ctx.accounts.token_program.to_account_info(),
                Transfer {
                    from: ctx.accounts.vault.to_account_info(),
                    to: ctx.accounts.borrower_token_account.to_account_info(),
                    authority: ctx.accounts.pool.to_account_info(),
                },
                signer_seeds,
            ),
            amount,
        )?;

        // Update loan counters
        let pool = &mut ctx.accounts.pool;
        pool.total_loans_executed = pool.total_loans_executed.checked_add(1).ok_or(FlashLoanError::MathOverflow)?;

        emit!(FlashLoanBorrowed {
            borrower: ctx.accounts.borrower.key(),
            mint: ctx.accounts.pool.mint,
            amount,
            fee,
            required_repay_amount,
        });

        Ok(())
    }

    /// Repay borrowed funds plus fee back to the vault
    pub fn repay_flash_loan(ctx: Context<RepayFlashLoan>, amount: u64) -> Result<()> {
        let fee = calculate_fee(amount)?;
        let total_repay_amount = amount.checked_add(fee).ok_or(FlashLoanError::MathOverflow)?;

        token::transfer(
            CpiContext::new(
                ctx.accounts.token_program.to_account_info(),
                Transfer {
                    from: ctx.accounts.borrower_token_account.to_account_info(),
                    to: ctx.accounts.vault.to_account_info(),
                    authority: ctx.accounts.borrower.to_account_info(),
                },
            ),
            total_repay_amount,
        )?;

        // Update fees collected
        let pool = &mut ctx.accounts.pool;
        pool.total_fees_collected = pool.total_fees_collected.checked_add(fee).ok_or(FlashLoanError::MathOverflow)?;

        emit!(FlashLoanRepaid {
            borrower: ctx.accounts.borrower.key(),
            mint: ctx.accounts.pool.mint,
            amount,
            fee,
            total_repay_amount,
        });

        Ok(())
    }

    /// Deposit liquidity into pool
    pub fn deposit_liquidity(ctx: Context<DepositLiquidity>, amount: u64) -> Result<()> {
        require!(amount > 0, FlashLoanError::ZeroDepositAmount);

        token::transfer(
            CpiContext::new(
                ctx.accounts.token_program.to_account_info(),
                Transfer {
                    from: ctx.accounts.depositor_token_account.to_account_info(),
                    to: ctx.accounts.vault.to_account_info(),
                    authority: ctx.accounts.depositor.to_account_info(),
                },
            ),
            amount,
        )?;

        emit!(LiquidityDeposited {
            depositor: ctx.accounts.depositor.key(),
            mint: ctx.accounts.pool.mint,
            amount,
        });

        Ok(())
    }

    /// Withdraw accumulated fees or liquidity (Admin only)
    pub fn withdraw_admin(ctx: Context<WithdrawAdmin>, amount: u64) -> Result<()> {
        require!(amount > 0, FlashLoanError::ZeroWithdrawAmount);
        require!(
            ctx.accounts.vault.amount >= amount,
            FlashLoanError::InsufficientVaultLiquidity
        );

        let mint_key = ctx.accounts.pool.mint;
        let pool_bump = ctx.accounts.pool.bump;
        let signer_seeds: &[&[&[u8]]] = &[&[
            POOL_SEED,
            mint_key.as_ref(),
            &[pool_bump],
        ]];

        token::transfer(
            CpiContext::new_with_signer(
                ctx.accounts.token_program.to_account_info(),
                Transfer {
                    from: ctx.accounts.vault.to_account_info(),
                    to: ctx.accounts.admin_token_account.to_account_info(),
                    authority: ctx.accounts.pool.to_account_info(),
                },
                signer_seeds,
            ),
            amount,
        )?;

        emit!(AdminWithdrawn {
            admin: ctx.accounts.admin.key(),
            mint: ctx.accounts.pool.mint,
            amount,
        });

        Ok(())
    }
}

// ----------------- Helper Functions -----------------
pub fn calculate_fee(amount: u64) -> Result<u64> {
    let fee = (amount as u128)
        .checked_mul(PROTOCOL_FEE_BPS as u128)
        .and_then(|v| v.checked_div(BPS_DENOMINATOR as u128))
        .ok_or(FlashLoanError::MathOverflow)?;
    
    // Minimum fee of 1 base unit if amount > 0
    let fee_u64 = if fee == 0 && amount > 0 { 1 } else { fee as u64 };
    Ok(fee_u64)
}

// ----------------- Account Contexts -----------------

#[derive(Accounts)]
pub struct InitializePool<'info> {
    #[account(mut)]
    pub admin: Signer<'info>,

    pub mint: Account<'info, Mint>,

    #[account(
        init,
        payer = admin,
        space = 8 + Pool::INIT_SPACE,
        seeds = [POOL_SEED, mint.key().as_ref()],
        bump
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        init,
        payer = admin,
        seeds = [VAULT_SEED, mint.key().as_ref()],
        bump,
        token::mint = mint,
        token::authority = pool,
    )]
    pub vault: Account<'info, TokenAccount>,

    pub token_program: Program<'info, Token>,
    pub system_program: Program<'info, System>,
    pub rent: Sysvar<'info, Rent>,
}

#[derive(Accounts)]
pub struct BorrowFlashLoan<'info> {
    #[account(mut)]
    pub borrower: Signer<'info>,

    #[account(
        mut,
        seeds = [POOL_SEED, pool.mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [VAULT_SEED, pool.mint.as_ref()],
        bump = pool.vault_bump,
    )]
    pub vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        constraint = borrower_token_account.mint == pool.mint @ FlashLoanError::InvalidMint,
    )]
    pub borrower_token_account: Account<'info, TokenAccount>,

    /// CHECK: Instructions Sysvar required for atomic transaction introspection
    #[account(address = INSTRUCTIONS_SYSVAR_ID @ FlashLoanError::InvalidInstructionsSysvar)]
    pub instructions_sysvar: AccountInfo<'info>,

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

#[derive(Accounts)]
pub struct RepayFlashLoan<'info> {
    #[account(mut)]
    pub borrower: Signer<'info>,

    #[account(
        mut,
        seeds = [POOL_SEED, pool.mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [VAULT_SEED, pool.mint.as_ref()],
        bump = pool.vault_bump,
    )]
    pub vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        constraint = borrower_token_account.mint == pool.mint @ FlashLoanError::InvalidMint,
    )]
    pub borrower_token_account: Account<'info, TokenAccount>,

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

#[derive(Accounts)]
pub struct DepositLiquidity<'info> {
    #[account(mut)]
    pub depositor: Signer<'info>,

    #[account(
        seeds = [POOL_SEED, pool.mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [VAULT_SEED, pool.mint.as_ref()],
        bump = pool.vault_bump,
    )]
    pub vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        constraint = depositor_token_account.mint == pool.mint @ FlashLoanError::InvalidMint,
    )]
    pub depositor_token_account: Account<'info, TokenAccount>,

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

#[derive(Accounts)]
pub struct WithdrawAdmin<'info> {
    #[account(
        mut,
        constraint = admin.key() == pool.admin @ FlashLoanError::UnauthorizedAdmin
    )]
    pub admin: Signer<'info>,

    #[account(
        seeds = [POOL_SEED, pool.mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [VAULT_SEED, pool.mint.as_ref()],
        bump = pool.vault_bump,
    )]
    pub vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        constraint = admin_token_account.mint == pool.mint @ FlashLoanError::InvalidMint,
    )]
    pub admin_token_account: Account<'info, TokenAccount>,

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

// ----------------- State -----------------

#[account]
#[derive(InitSpace)]
pub struct Pool {
    pub admin: Pubkey,
    pub mint: Pubkey,
    pub vault: Pubkey,
    pub total_fees_collected: u64,
    pub total_loans_executed: u64,
    pub bump: u8,
    pub vault_bump: u8,
}

// ----------------- Events -----------------

#[event]
pub struct PoolInitialized {
    pub admin: Pubkey,
    pub mint: Pubkey,
    pub vault: Pubkey,
}

#[event]
pub struct FlashLoanBorrowed {
    pub borrower: Pubkey,
    pub mint: Pubkey,
    pub amount: u64,
    pub fee: u64,
    pub required_repay_amount: u64,
}

#[event]
pub struct FlashLoanRepaid {
    pub borrower: Pubkey,
    pub mint: Pubkey,
    pub amount: u64,
    pub fee: u64,
    pub total_repay_amount: u64,
}

#[event]
pub struct LiquidityDeposited {
    pub depositor: Pubkey,
    pub mint: Pubkey,
    pub amount: u64,
}

#[event]
pub struct AdminWithdrawn {
    pub admin: Pubkey,
    pub mint: Pubkey,
    pub amount: u64,
}

// ----------------- Errors -----------------

#[error_code]
pub enum FlashLoanError {
    #[msg("Borrow amount must be greater than zero.")]
    ZeroBorrowAmount,
    #[msg("Deposit amount must be greater than zero.")]
    ZeroDepositAmount,
    #[msg("Withdraw amount must be greater than zero.")]
    ZeroWithdrawAmount,
    #[msg("Insufficient liquidity available in the vault.")]
    InsufficientVaultLiquidity,
    #[msg("Repayment instruction is missing from this transaction.")]
    RepaymentInstructionMissing,
    #[msg("Mathematical overflow occurred.")]
    MathOverflow,
    #[msg("Token account mint does not match pool mint.")]
    InvalidMint,
    #[msg("Invalid instructions sysvar account.")]
    InvalidInstructionsSysvar,
    #[msg("Unauthorized: only the pool admin can withdraw.")]
    UnauthorizedAdmin,
}