lib.rs

#![allow(unexpected_cfgs)]
use anchor_lang::{
    prelude::*,
    solana_program::{instruction::Instruction, program::invoke, system_instruction},
};
use anchor_spl::{
    token::spl_token,
};

mod events;
mod state;
mod contexts;

use contexts::*;
use events::*;
use state::*;

declare_id!("EeT88L4VvDieXZgko1LuG1UYLc2jGaqy5Wichda7EfG5");

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

    pub fn initialize(ctx: Context<Initialize>) -> Result<()> {

        ctx.accounts.vault.set_inner(Vault::new(ctx.accounts.caller.key()));

        emit!(VaultInitialized {
            authority: ctx.accounts.caller.key(),
            vault: ctx.accounts.vault.key(),
        });

        let vault = &mut ctx.accounts.vault;
        msg!("vault amount to buy = {}", vault.next_round_buy_amount);

        Ok(())
    }

    pub fn set_executor(ctx: Context<SetExecutor>, new_executor: Option<Pubkey>) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        let caller = &ctx.accounts.caller.key();

        require_authority_role(caller, vault)?;

        vault.executor = new_executor;

        emit!(ExecutorChanged {
            new_executor,
            authority: *caller,
        });

        Ok(())
    }

    pub fn set_user_investment(
        ctx: Context<SetUserInvestment>,
        user: Pubkey,
        investment_per_token: u64,
        is_active: bool,
    ) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        let caller = &ctx.accounts.caller.key();

        require_executor_role(caller, vault)?;

        let user_account = &mut ctx.accounts.user_account;
        require!(user_account.user == user, VaultErrorCode::UserMismatch);

        // === CALCULATE OLD AND NEW CONTRIBUTION ===
        let old_contribution = if user_account.sol_balance >= user_account.investment_per_token
                                && user_account.is_active {
            user_account.investment_per_token
        } else {
            0
        };

        let new_contribution = if user_account.sol_balance >= investment_per_token && is_active {
            investment_per_token
        } else {
            0
        };

        // === UPDATE USER SETTINGS ===
        user_account.investment_per_token = investment_per_token;
        user_account.is_active = is_active;

        // === UPDATE VAULT BUY AMOUNT ===
        // Only update during Deposit stage
        if vault.stage == VaultStage::Deposit {
            vault.update_next_round_amount(old_contribution, new_contribution)?;
        }

        emit!(UserSettingsUpdated {
            user,
            investment_per_token,
            is_active,
        });

        Ok(())
    }

    pub fn deposit(
        ctx: Context<Deposit>,
        user: Pubkey,
    ) -> Result<()> {
        // === SETUP ===
        let vault = &mut ctx.accounts.vault;
        let user_account = &mut ctx.accounts.user_account;
        let caller = &ctx.accounts.caller.key();

        // === VALIDATION ===
        require!(vault.stage == VaultStage::Deposit, VaultErrorCode::WrongStage);
        require_executor_role(caller, vault)?;

        // === ACCOUNT INITIALIZATION (this will deduct rent automatically) ===
        if user_account.is_uninitialized() {
            user_account.set_inner(UserAccount::new(user));
        }

        // === CALCULATE DEPOSIT AMOUNT AFTER ACCOUNT CREATION ===
        // Now check what's actually left after account creation
        let deposit_amount = ctx.accounts.deposit_account.lamports();
        require!(deposit_amount > 0, VaultErrorCode::InsufficientBalance);

        // === CALCULATE & UPDATE BUY AMOUNT ===
        // If user becomes eligible to participate, add to buy amount
        if user_account.sol_balance >= user_account.investment_per_token && user_account.is_active {    // TODO: consider removing active/inactive
            vault.add_to_next_round(user_account.investment_per_token)?;
        }

        // === BALANCE UPDATES ===
        user_account.add_sol_balance(deposit_amount)?;
        vault.add_deposits(deposit_amount)?;

        // === SOL TRANSFER ===
        let transfer_instruction = system_instruction::transfer(
            &ctx.accounts.deposit_account.key(),
            &vault.key(),
            deposit_amount,
        );

        invoke(
            &transfer_instruction,
            &[
                ctx.accounts.deposit_account.to_account_info(),
                vault.to_account_info(),
                ctx.accounts.system_program.to_account_info(),
            ],
        )?;

        emit!(DepositMade {
            user,
            amount: deposit_amount,
            new_user_balance: user_account.sol_balance,
            total_vault_deposits: vault.deposits,
            total_buy_amount: vault.next_round_buy_amount,
        });

        Ok(())
    }

    pub fn buy_tokens_collective(
        ctx: Context<BuyTokensCollective>,
        target_token: Pubkey,
        jupiter_swap_data: Vec<u8>,
        min_amount_out: u64,
    ) -> Result<()> {
        // === SETUP & VALIDATION ===
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        let buy_amount = vault.next_round_buy_amount;
        require!(buy_amount > 0, VaultErrorCode::InsufficientFunds);

        msg!("Instant buy for {} with {} SOL", target_token, buy_amount);

        require!(vault.stage == VaultStage::Buy, VaultErrorCode::WrongStage);

        // === TRACK STARTING BALANCE ===
        let caller_token_balance_before = ctx.accounts.authority_token_account.amount; // TODO: consider changing the name of the token account
        msg!("Authority token balance before: {}", caller_token_balance_before);

        // === EXECUTE JUPITER SWAP ===
        // Preparing account metadata: Jupiter's swap instruction expects accounts in a specific format (Vec<AccountMeta>),
        // but Anchor gives you AccountInfo. This code converts between the two formats.
        // remaining_accounts -> "all the other accounts Jupiter needs" (token mints, ATAs, market accounts, etc).
        let accounts: Vec<AccountMeta> = ctx
            .remaining_accounts
            .iter()
            .map(|acc| {
                AccountMeta {
                    pubkey: *acc.key,                                   // The account's address
                    is_signer: acc.key == &ctx.accounts.caller.key(),   // Is this account signing?
                    is_writable: acc.is_writable,                       // Can this account be modified?
                }
            })
            .collect();

        // TODO: test with let account_infos: Vec<AccountInfo> = ctx.remaining_accounts.to_vec();
        // or with invoke(&instruction, ctx.remaining_accounts)?;
        let account_infos: Vec<AccountInfo> = ctx
            .remaining_accounts
            .iter()
            .map(|acc| AccountInfo { ..acc.clone() })
            .collect();

        invoke(
            &Instruction {
                program_id: ctx.accounts.jupiter_program.key(),
                accounts,
                data: jupiter_swap_data,
            },
            &account_infos,
        )?;

        msg!("Jupiter swap completed");

        // === CALCULATE & VALIDATE TOKENS RECEIVED ===
        ctx.accounts.authority_token_account.reload()?;
        let caller_token_balance_after = ctx.accounts.authority_token_account.amount;
        let tokens_received = caller_token_balance_after
            .checked_sub(caller_token_balance_before)
            .ok_or(VaultErrorCode::InvalidCalculation)?;

        require!(
            tokens_received >= min_amount_out,
            VaultErrorCode::InsufficientTokensReceived
        );
        msg!("Tokens received by authority: {}", tokens_received);

        // === TRANSFER TOKENS TO VAULT ===
        if tokens_received > 0 {
            let transfer_cpi_accounts = anchor_spl::token::Transfer {
                from: ctx.accounts.authority_token_account.to_account_info(),
                to: ctx.accounts.vault_token_account.to_account_info(),
                authority: ctx.accounts.caller.to_account_info(),
            };
            let transfer_cpi_program = ctx.accounts.output_token_program.to_account_info();
            let transfer_cpi_ctx = CpiContext::new(
                transfer_cpi_program,
                transfer_cpi_accounts
            );

            anchor_spl::token::transfer(transfer_cpi_ctx, tokens_received)?;
            msg!("Transferred {} tokens from authority to vault", tokens_received);
        }

        // === CLEANUP WSOL: Handle any remaining WSOL ===
        ctx.accounts.user_wsol_token_account.reload()?;
        let remaining_wsol = ctx.accounts.user_wsol_token_account.amount;

        if remaining_wsol > 0 {
            msg!("Closing WSOL account with {} remaining", remaining_wsol);

            let close_ix = spl_token::instruction::close_account(
                &spl_token::ID,
                &ctx.accounts.user_wsol_token_account.key(),
                &ctx.accounts.caller.key(),
                &ctx.accounts.caller.key(),
                &[],
            )?;

            invoke(
                &close_ix,
                &[
                    ctx.accounts.user_wsol_token_account.to_account_info(),
                    ctx.accounts.caller.to_account_info(),
                    ctx.accounts.caller.to_account_info(),
                ],
            )?;

            msg!("WSOL account closed, SOL returned to caller");
        }

        // === UPDATE VAULT STATE ===
        vault.add_or_update_token_balance(target_token, tokens_received)?;

        emit!(CollectiveBuyCompleted {
            swap_amount: 1u64,  // TODO: fix this
            target_token,
            tokens_received,
            effective_price: if tokens_received > 0 {
                (buy_amount as u128 * 1_000_000) / tokens_received as u128
            } else { 0 },
        });

        Ok(())
    }

    pub fn sell_tokens(
        ctx: Context<SellTokens>,
        token_mint: Pubkey,
        token_amount_to_sell: u64,
        jupiter_swap_data: Vec<u8>,
        min_sol_out: u64,
    ) -> Result<()> {
        // === SETUP & VALIDATION ===
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;

        require!(vault.stage == VaultStage::Buy, VaultErrorCode::WrongStage);

        // Check vault has enough tokens to sell
        let vault_token_balance = vault.get_token_balance(token_mint)?;
        require!(
            vault_token_balance >= token_amount_to_sell,
            VaultErrorCode::InsufficientTokenBalance
        );
        require!(token_amount_to_sell > 0, VaultErrorCode::InsufficientFunds); // TODO; change this error

        msg!("Collective sell: {} tokens of mint {}", token_amount_to_sell, token_mint);

        // === TRACK STARTING BALANCE ===
        let authority_wsol_balance_before = ctx.accounts.authority_wsol_account.amount;
        msg!("Authority WSOL balance before: {}", authority_wsol_balance_before);

        // === PREPARE TOKENS FOR SWAP ===
        // Transfer tokens from vault to authority for the swap
        let vault_seeds = &[VAULT_SEED, &[ctx.bumps.vault]];
        let signer_seeds = &[&vault_seeds[..]];

        let transfer_to_authority_accounts = anchor_spl::token::Transfer {
            from: ctx.accounts.vault_token_account.to_account_info(),
            to: ctx.accounts.authority_token_account.to_account_info(),
            authority: vault.to_account_info(),
        };
        let transfer_to_authority_ctx = CpiContext::new_with_signer(
            ctx.accounts.input_token_program.to_account_info(),
            transfer_to_authority_accounts,
            signer_seeds,
        );

        anchor_spl::token::transfer(transfer_to_authority_ctx, token_amount_to_sell)?;
        msg!("Transferred {} tokens from vault to authority for swap", token_amount_to_sell);

        // === EXECUTE JUPITER SWAP ===
        let accounts: Vec<AccountMeta> = ctx
            .remaining_accounts
            .iter()
            .map(|acc| {
                AccountMeta {
                    pubkey: *acc.key,
                    is_signer: acc.key == &ctx.accounts.caller.key(),
                    is_writable: acc.is_writable,
                }
            })
            .collect();

        let account_infos: Vec<AccountInfo> = ctx
            .remaining_accounts
            .iter()
            .map(|acc| AccountInfo { ..acc.clone() })
            .collect();

        invoke(
            &Instruction {
                program_id: ctx.accounts.jupiter_program.key(),
                accounts,
                data: jupiter_swap_data,
            },
            &account_infos,
        )?;

        msg!("Jupiter swap completed");

        // === CALCULATE & VALIDATE SOL RECEIVED ===
        ctx.accounts.authority_wsol_account.reload()?;
        let authority_wsol_balance_after = ctx.accounts.authority_wsol_account.amount;
        let wsol_received = authority_wsol_balance_after
            .checked_sub(authority_wsol_balance_before)
            .ok_or(VaultErrorCode::InvalidCalculation)?;

        require!(
            wsol_received >= min_sol_out,
            VaultErrorCode::InsufficientTokensReceived
        );
        msg!("WSOL received from swap: {}", wsol_received);

        // === CLEANUP WSOL: Convert back to SOL and send to vault ===
        if wsol_received > 0 {
            let close_ix = spl_token::instruction::close_account(
                &spl_token::ID,
                &ctx.accounts.authority_wsol_account.key(),
                &vault.key(), // SOL goes to vault
                &ctx.accounts.caller.key(),
                &[],
            )?;

            invoke(
                &close_ix,
                &[
                    ctx.accounts.authority_wsol_account.to_account_info(),
                    vault.to_account_info(), // SOL recipient is the vault
                    ctx.accounts.caller.to_account_info(),
                ],
            )?;

            msg!("WSOL account closed, {} SOL sent to vault", wsol_received);
        }

        // === UPDATE VAULT STATE ===
        vault.subtract_token_balance(token_mint, token_amount_to_sell)?;

        emit!(CollectiveSellCompleted {
            token_mint,
            tokens_sold: token_amount_to_sell,
            sol_received: wsol_received,
            effective_price: if token_amount_to_sell > 0 {
                (wsol_received as u128 * 1_000_000) / token_amount_to_sell as u128
            } else { 0 },
        });

        Ok(())
    }

    pub fn emergency_withdraw_tokens(
        ctx: Context<EmergencyWithdrawTokens>,
        token_mint: Pubkey,
    ) -> Result<()> {
        let vault = &ctx.accounts.vault;

        // Only vault authority can call this
        require_authority_role(&ctx.accounts.authority.key(), vault)?;

        let token_balance = ctx.accounts.vault_token_account.amount;

        if token_balance > 0 {
            // Create signer seeds for vault PDA
            let vault_seeds = &[VAULT_SEED, &[ctx.bumps.vault]];
            let signer_seeds = &[&vault_seeds[..]];

            // Transfer all tokens from vault to authority
            let cpi_accounts = anchor_spl::token::Transfer {
                from: ctx.accounts.vault_token_account.to_account_info(),
                to: ctx.accounts.authority_token_account.to_account_info(),
                authority: vault.to_account_info(),
            };

            let cpi_ctx = CpiContext::new_with_signer(
                ctx.accounts.token_program.to_account_info(),
                cpi_accounts,
                signer_seeds,
            );

            anchor_spl::token::transfer(cpi_ctx, token_balance)?;

            msg!("🚨 Emergency: Withdrew {} tokens of mint {} to authority",
                token_balance, token_mint);
        }

        Ok(())
    }

    pub fn emergency_withdraw_sol(
        ctx: Context<EmergencyWithdrawSol>,
        amount: u64,
    ) -> Result<()> {
        let vault = &mut ctx.accounts.vault;

        require_authority_role(&ctx.accounts.authority.key(), vault)?;

        let vault_balance = vault.to_account_info().lamports();
        let withdraw_amount = if amount == 0 {
            vault_balance.saturating_sub(5_000_000) // Leave 0.005 SOL for rent
        } else {
            amount
        };

        require!(vault_balance >= withdraw_amount, VaultErrorCode::InsufficientFunds);

        // Transfer SOL from vault to authority
        **vault.to_account_info().try_borrow_mut_lamports()? -= withdraw_amount;
        **ctx.accounts.authority.to_account_info().try_borrow_mut_lamports()? += withdraw_amount;

        msg!("🚨 Emergency: Withdrew {} SOL from vault to authority",
            withdraw_amount as f64 / 1_000_000_000.0);

        Ok(())
    }

    pub fn start_buy_stage(ctx: Context<StartBuyStage>) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require_stage(vault, VaultStage::Deposit)?;
        require!(vault.next_round_buy_amount > 0, VaultErrorCode::SetupNotCompleted);

        // ✅ Calculate fees here (only once when entering Buy stage)
        let total_investment = vault.next_round_buy_amount;
        let fee_amount = (total_investment * 2) / 100;
        let net_buy_amount = total_investment - fee_amount;  // TODO: check this, the fee was already applied?

        // ✅ Track fees and update buy amount
        vault.add_fees(fee_amount)?;
        vault.set_current_spend(net_buy_amount);

        // ✅ Change stage
        vault.stage = VaultStage::Buy;

        emit!(SolFeesCollected {
            fee_amount,
            total_contribution: total_investment,
            net_buy_amount,
            fee_percentage: 2,
            round: vault.current_round,
        });

        emit!(StageChanged {
            old_stage: VaultStage::Deposit,
            new_stage: VaultStage::Buy,
            round: vault.current_round,
            total_buy_amount: vault.next_round_buy_amount,
        });

        Ok(())
    }

    pub fn start_update_stage(ctx: Context<StartUpdateStage>) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require!(vault.stage == VaultStage::Buy, VaultErrorCode::WrongStage);

        vault.stage = VaultStage::Update;

        emit!(StageChanged {
            old_stage: VaultStage::Buy,
            new_stage: VaultStage::Update,
            round: vault.current_round,
            total_buy_amount: vault.next_round_buy_amount,
        });

        Ok(())
    }

    // === O(1) ALTERNATIVE: Just store round metadata ===
    pub fn finalize_round_o1(
        ctx: Context<UpdateUserBalancesBatch>,
        token_mint: Pubkey,
        total_tokens_received: u64,
    ) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require_stage(vault, VaultStage::Update)?;

        let total_sol_spent = vault.current_round_spend;
        require!(total_sol_spent > 0, VaultErrorCode::InsufficientFunds);

        // ✅ O(1): Just store the round snapshot - no user iteration
        vault.round_snapshots.push(RoundSnapshot {
            round: vault.current_round,
            total_sol_spent,
            total_tokens_received,
            token_mint,
        });

        // Update vault's token holdings
        vault.add_or_update_token_balance(token_mint, total_tokens_received)?;

        emit!(RoundFinalized {
            round: vault.current_round,
            token_mint,
            total_sol_spent,
            total_tokens_received,
        });

        Ok(())
    }

    pub fn update_user_balances_batch(
        ctx: Context<UpdateUserBalancesBatch>,
        token_mint: Pubkey,
        total_tokens_received: u64,
        users: Vec<Pubkey>,
    ) -> Result<()> {
        // === VALIDATION ===
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require_stage(vault, VaultStage::Update)?;
        require!(!users.is_empty(), VaultErrorCode::EmptyBatch);

        let total_sol_spent = vault.current_round_spend;
        require!(total_sol_spent > 0, VaultErrorCode::InsufficientFunds);

        // === PROCESS ALL USERS ===
        for (account_index, user_pubkey) in users.iter().enumerate() {
            let user_account_info = &ctx.remaining_accounts[account_index];

            modify_account::<UserAccount, _>(user_account_info, |user_account| {
                require!(user_account.user == *user_pubkey, VaultErrorCode::UserMismatch);

                let contribution = calculate_user_contribution(user_account);
                if contribution > 0 {

                    let tokens = calculate_token_allocation(contribution, total_tokens_received, total_sol_spent);
                    user_account.subtract_sol_balance(contribution)?;
                    user_account.add_or_update_token_balance(token_mint, tokens)?;
                }

                // === NEXT ROUND: Calculate what user will contribute ===
                // After the SOL deduction above, check if they can still participate next round
                let next_round_contribution = calculate_user_contribution(user_account);

                if next_round_contribution > 0 {
                    // Add to next round's buy amount
                    vault.add_to_next_round(next_round_contribution)?;
                }

                Ok(())
            })?;
        }

        // === EMIT SINGLE BATCH EVENT ===
        emit!(BatchProcessed {
            users_count: users.len() as u32,
            total_sol_spent,
            total_tokens_distributed: total_tokens_received,
            round: vault.current_round,
        });

        // === EMIT NEXT ROUND PREPARATION EVENT ===
        // emit!(NextRoundPrepared {
        //     next_round: vault.current_round + 1,
        //     estimated_buy_amount: vault.next_round_buy_amount,
        //     users_processed: users.len() as u32,
        // });

        Ok(())
    }

    pub fn start_withdraw_stage(ctx: Context<StartWithdrawStage>) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require!(vault.stage == VaultStage::Update, VaultErrorCode::WrongStage);

        vault.stage = VaultStage::Withdraw;

        emit!(StageChanged {
            old_stage: VaultStage::Update,
            new_stage: VaultStage::Withdraw,
            round: vault.current_round,
            total_buy_amount: vault.next_round_buy_amount,
        });

        Ok(())
    }

    // TODO: deal with the reset of fees between rounds, in update balances function
    pub fn start_new_round(ctx: Context<StartNewRound>) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require!(vault.stage == VaultStage::Withdraw, VaultErrorCode::WrongStage);

        vault.stage = VaultStage::Deposit;  // Back to Deposit for new round
        vault.current_round += 1;
        vault.accumulated_fees = 0;

        emit!(NewRoundStarted {
            round: vault.current_round,
        });

        emit!(StageChanged {
            old_stage: VaultStage::Withdraw,
            new_stage: VaultStage::Deposit,
            round: vault.current_round,
            total_buy_amount: vault.next_round_buy_amount,
        });

        Ok(())
    }

    pub fn withdraw_sol(
        ctx: Context<WithdrawSol>,
        user: Pubkey,
        destination: Pubkey,  // ← Add destination parameter
        amount: u64, // 0 means withdraw all
    ) -> Result<()> {
        let vault = &mut ctx.accounts.vault;
        let user_account = &mut ctx.accounts.user_account;

        // === VALIDATION ===
        require!(vault.stage == VaultStage::Withdraw, VaultErrorCode::WrongStage);
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require!(user_account.user == user, VaultErrorCode::UserMismatch);

        // Calculate withdrawal amount
        let withdraw_amount = if amount == 0 {
            user_account.sol_balance // Withdraw all
        } else {
            amount
        };

        require!(withdraw_amount > 0, VaultErrorCode::InsufficientFunds);
        require!(user_account.sol_balance >= withdraw_amount, VaultErrorCode::InsufficientBalance);

        // Check vault has enough SOL
        let vault_balance = vault.to_account_info().lamports();
        require!(vault_balance >= withdraw_amount, VaultErrorCode::InsufficientVaultBalance);

        // Validate destination matches the account provided
        require!(
            destination == ctx.accounts.destination.key(),
            VaultErrorCode::InvalidDestination
        );

        // === UPDATE BALANCES ===
        user_account.subtract_sol_balance(withdraw_amount)?;
        vault.deposits = vault.deposits
            .checked_sub(withdraw_amount)
            .ok_or(VaultErrorCode::MathOverflow)?;

        // === TRANSFER SOL ===
        // Transfer SOL from vault to destination
        **vault.to_account_info().try_borrow_mut_lamports()? -= withdraw_amount;
        **ctx.accounts.destination.to_account_info().try_borrow_mut_lamports()? += withdraw_amount;

        emit!(WithdrawalCompleted {
            user,
            asset_type: "SOL".to_string(),
            amount: withdraw_amount,
            destination,
            new_user_sol_balance: user_account.sol_balance,
            remaining_vault_deposits: vault.deposits,
        });

        Ok(())
    }

    pub fn withdraw_tokens(
        ctx: Context<WithdrawTokens>,
        user: Pubkey,
        token_mint: Pubkey,
        destination_authority: Pubkey,  // ← Add destination authority parameter
        amount: u64, // 0 means withdraw all
    ) -> Result<()> {
        let vault = &ctx.accounts.vault;
        let user_account = &mut ctx.accounts.user_account;

                // === VALIDATION ===
        require!(vault.stage == VaultStage::Withdraw, VaultErrorCode::WrongStage);
        require_executor_role(&ctx.accounts.caller.key(), vault)?;
        require!(user_account.user == user, VaultErrorCode::UserMismatch);

        // Validate destination authority matches the account provided
        require!(
            destination_authority == ctx.accounts.destination_authority.key(),
            VaultErrorCode::InvalidDestination
        );

        // ✅ LAZY EVALUATION: Update user balances only when needed
        user_account.apply_pending_updates(vault, vault.current_round)?;

        // Get user's token balance (now includes pending tokens)
        let user_token_balance = user_account.get_token_balance(token_mint)?;
        require!(user_token_balance > 0, VaultErrorCode::InsufficientTokenBalance);

        // Calculate withdrawal amount
        let withdraw_amount = if amount == 0 {
            user_token_balance // Withdraw all
        } else {
            amount
        };

        require!(withdraw_amount > 0, VaultErrorCode::InsufficientFunds);
        require!(user_token_balance >= withdraw_amount, VaultErrorCode::InsufficientTokenBalance);

        // Check vault has enough tokens
        let vault_token_balance = ctx.accounts.vault_token_account.amount;
        require!(vault_token_balance >= withdraw_amount, VaultErrorCode::InsufficientTokenBalance);

        // === UPDATE USER BALANCE ===
        user_account.subtract_token_balance(token_mint, withdraw_amount)?;

        // === TRANSFER TOKENS ===
        // Create signer seeds for vault PDA
        let vault_seeds = &[VAULT_SEED, &[ctx.bumps.vault]];
        let signer_seeds = &[&vault_seeds[..]];

        let transfer_accounts = anchor_spl::token::Transfer {
            from: ctx.accounts.vault_token_account.to_account_info(),
            to: ctx.accounts.destination_token_account.to_account_info(),
            authority: vault.to_account_info(),
        };

        let transfer_ctx = CpiContext::new_with_signer(
            ctx.accounts.token_program.to_account_info(),
            transfer_accounts,
            signer_seeds,
        );

        anchor_spl::token::transfer(transfer_ctx, withdraw_amount)?;

        emit!(WithdrawalCompleted {
            user,
            asset_type: format!("TOKEN:{}", token_mint),
            amount: withdraw_amount,
            destination: ctx.accounts.destination_token_account.key(),
            new_user_sol_balance: user_account.sol_balance, // SOL balance unchanged
            remaining_vault_deposits: vault.deposits, // Deposits unchanged for token withdrawal
        });

        Ok(())
    }

    pub fn withdraw_sol_fees(
        ctx: Context<WithdrawSolFees>,
        amount: u64, // 0 means withdraw all fees
    ) -> Result<()> {
        let vault = &mut ctx.accounts.vault;

        // Only authority can withdraw fees
        require_authority_role(&ctx.accounts.authority.key(), vault)?;

        let withdraw_amount = if amount == 0 {
            vault.accumulated_fees // Withdraw all fees
        } else {
            amount
        };

        require!(withdraw_amount > 0, VaultErrorCode::InsufficientFunds);
        require!(vault.accumulated_fees >= withdraw_amount, VaultErrorCode::InsufficientFunds);

        // Check vault has enough SOL
        let vault_balance = vault.to_account_info().lamports();
        require!(vault_balance >= withdraw_amount, VaultErrorCode::InsufficientVaultBalance);

        // Update fee tracking
        vault.subtract_fees(withdraw_amount)?;

        // Transfer SOL from vault to authority
        **vault.to_account_info().try_borrow_mut_lamports()? -= withdraw_amount;
        **ctx.accounts.authority.to_account_info().try_borrow_mut_lamports()? += withdraw_amount;

        emit!(SolFeesWithdrawn {
            amount: withdraw_amount,
            remaining_fee_balance: vault.accumulated_fees,
            authority: ctx.accounts.authority.key(),
        });

        Ok(())
    }

}

// Only need these two helper functions:
fn calculate_user_contribution(user_account: &UserAccount) -> u64 {
    if user_account.sol_balance >= user_account.investment_per_token && user_account.is_active {
        user_account.investment_per_token
    } else {
        0
    }
}

fn calculate_token_allocation(user_sol_contribution: u64, total_tokens_received: u64, total_sol_spent: u64) -> u64 {
    if user_sol_contribution == 0 {
        return 0;
    }

    // Calculate user's net contribution (after 2% fee)
    let user_fee = (user_sol_contribution * 2) / 100;
    let user_net_contribution = user_sol_contribution - user_fee;

    // Calculate proportional share based on net contributions
    // Formula: (user_net_contribution / total_net_spent) * total_tokens
    (user_net_contribution as u128 * total_tokens_received as u128 / total_sol_spent as u128) as u64
}

MATH_ANALYSIS.md

Overflow and Underflow Analysis for Vault State Variables

Overview

This document analyzes all state variables in the vault program for potential overflow and underflow vulnerabilities. The analysis covers the Vault account, UserAccount account, and related data structures.

State Variable Analysis

1. Vault Account State Variables

1.1 deposits: u64

Range: 0 to 18,446,744,073,709,551,615 lamports (~18.4 quintillion lamports or ~18.4 billion SOL)

Overflow Scenarios:

• During deposit operations: Multiple large deposits could cause cumulative overflow
  • Location: vault.add_deposits(deposit_amount) in deposit() function
  • Mitigation: Uses checked_add() with proper error handling
  • Risk: LOW - Protected by checked arithmetic

Underflow Scenarios:

• During SOL withdrawals: Withdrawing more than total deposits
  • Location: vault.deposits.checked_sub(withdraw_amount) in withdraw_sol()
  • Mitigation: Uses checked_sub() with proper error handling
  • Risk: LOW - Protected by checked arithmetic

1.2 next_round_buy_amount: u64

Range: 0 to 18,446,744,073,709,551,615 lamports

Overflow Scenarios:

• During user investment updates: Accumulating investment amounts from many users

  • Location: vault.add_to_next_round(amount) in multiple functions
  • Mitigation: Uses checked_add() in add_to_next_round()
  • Risk: LOW - Protected by checked arithmetic

• During deposit operations: Adding user investment per token amounts

  • Location: vault.add_to_next_round(user_account.investment_per_token) in deposit()
  • Mitigation: Uses checked_add() through add_to_next_round()
  • Risk: LOW - Protected by checked arithmetic

Underflow Scenarios:

• During investment amount adjustments: Reducing investment when users become inactive
  • Location: vault.subtract_from_next_round(amount) in update_next_round_amount()
  • Mitigation: Uses checked_sub() in subtract_from_next_round()
  • Risk: LOW - Protected by checked arithmetic

1.3 current_round_spend: u64

Range: 0 to 18,446,744,073,709,551,615 lamports

Overflow Scenarios:

• During direct assignment: No overflow risk as it's set directly
  • Location: vault.set_current_spend(net_buy_amount) in start_buy_stage()
  • Risk: NONE - Direct assignment, no arithmetic operations

Underflow Scenarios:

• No underflow risk: This variable is only set, never decremented
  • Risk: NONE

1.4 current_round: u64

Range: 0 to 18,446,744,073,709,551,615

Overflow Scenarios:

• During round increments: After ~18.4 quintillion rounds
  • Location: vault.current_round += 1 in start_new_round()
  • Mitigation: NONE - Uses unchecked increment
  • Risk: VERY LOW - Would require trillions of years to overflow

Underflow Scenarios:

• No underflow risk: This variable is only incremented, never decremented
  • Risk: NONE

1.5 accumulated_fees: u64

Range: 0 to 18,446,744,073,709,551,615 lamports

Overflow Scenarios:

• During fee accumulation: Adding 2% fees from each round
  • Location: vault.add_fees(fee_amount) in start_buy_stage()
  • Mitigation: Uses checked_add() in add_fees()
  • Risk: LOW - Protected by checked arithmetic

Underflow Scenarios:

• During fee withdrawals: Withdrawing more fees than accumulated
  • Location: vault.subtract_fees(amount) in withdraw_sol_fees()
  • Mitigation: Balance check before subtraction + manual subtraction
  • Risk: LOW - Protected by balance validation

1.6 token_holdings: Vec<TokenBalance>

TokenBalance.amount: u64 per token

Overflow Scenarios:

• During token purchases: Adding received tokens to existing balance
  • Location: vault.add_or_update_token_balance(target_token, tokens_received) in buy_tokens_collective()
  • Mitigation: Uses checked_add() in add_or_update_token_balance()
  • Risk: LOW - Protected by checked arithmetic

Underflow Scenarios:

• During token sales: Selling more tokens than vault holds
  • Location: vault.subtract_token_balance(token_mint, token_amount_to_sell) in sell_tokens()
  • Mitigation: Balance validation before subtraction
  • Risk: LOW - Protected by balance checks

2. UserAccount State Variables

2.1 sol_balance: u64

Range: 0 to 18,446,744,073,709,551,615 lamports

Overflow Scenarios:

• During deposits: Adding large deposit amounts
  • Location: user_account.add_sol_balance(deposit_amount) in deposit()
  • Mitigation: Uses checked_add() in add_sol_balance()
  • Risk: LOW - Protected by checked arithmetic

Underflow Scenarios:

• During withdrawals: Withdrawing more SOL than user balance

  • Location: user_account.subtract_sol_balance(withdraw_amount) in withdraw_sol()
  • Mitigation: Balance validation before subtraction
  • Risk: LOW - Protected by balance checks

• During balance updates: Deducting investment amounts during round processing

  • Location: user_account.subtract_sol_balance(contribution) in update_user_balances_batch()
  • Mitigation: Balance validation before subtraction
  • Risk: LOW - Protected by balance checks

2.2 investment_per_token: u64

Range: 0 to 18,446,744,073,709,551,615 lamports

Overflow Scenarios:

• During investment setting: Setting extremely large investment amounts
  • Location: Direct assignment in set_user_investment()
  • Risk: NONE - Direct assignment, no arithmetic operations

Underflow Scenarios:

• No underflow risk: This variable is only set, never decremented
  • Risk: NONE

2.3 token_balances: Vec<TokenBalance>

TokenBalance.amount: u64 per token

Overflow Scenarios:

• During token allocation: Adding tokens received from collective purchases
  • Location: user_account.add_or_update_token_balance(token_mint, tokens) in update_user_balances_batch()
  • Mitigation: Uses checked_add() in add_or_update_token_balance()
  • Risk: LOW - Protected by checked arithmetic

Underflow Scenarios:

• During token withdrawals: Withdrawing more tokens than user owns
  • Location: user_account.subtract_token_balance(token_mint, withdraw_amount) in withdraw_tokens()
  • Mitigation: Balance validation before subtraction
  • Risk: LOW - Protected by balance checks

3. Data Structure Variables

3.1 UserPortfolio (Event/View Structure)

• sol_balance: u64 - Same risks as UserAccount.sol_balance
• amount_per_round: u64 - Same risks as UserAccount.investment_per_token
• last_participated_round: u64 - Increment-only, very low overflow risk
• total_rounds_participated: u32 - Range: 0 to 4,294,967,295, increment-only

Overflow Scenarios:

• total_rounds_participated: After ~4.3 billion rounds per user
  • Risk: VERY LOW - Would require millions of years of participation

3.2 UserBuyData (Instruction Parameter)

• amount_spent: u64 - Externally provided, validated against user balance
• tokens_received: u64 - Externally provided, validated against actual tokens

Overflow Scenarios:

• Malicious input: External data could contain max u64 values
  • Risk: MEDIUM - Depends on proper validation in instruction handlers

Critical Mathematical Operations

4.1 Fee Calculations

let fee_amount = (total_investment * 2) / 100;

Overflow Scenarios:

• When total_investment > u64::MAX / 2 (extremely large)
• Risk: LOW - Would require ~9 billion SOL investment

4.2 Token Allocation Calculations

(user_net_contribution as u128 * total_tokens_received as u128 / total_sol_spent as u128) as u64

Overflow Scenarios:

• Intermediate calculation uses u128, preventing overflow in multiplication
• Final cast to u64 could truncate if result > u64::MAX
• Risk: LOW - Uses u128 for intermediate calculations

4.3 Price Calculations

(buy_amount as u128 * 1_000_000) / tokens_received as u128

Overflow Scenarios:

• Uses u128 for calculation, preventing overflow
• Risk: NONE - Properly protected

Recommendations

1. Add checked arithmetic to round increments:

   vault.current_round = vault.current_round.checked_add(1).ok_or(VaultErrorCode::MathOverflow)?;

2. Validate external inputs: Ensure all externally provided amounts are validated against available balances

3. Consider upper bounds: Implement reasonable maximum limits for investment amounts and deposits

4. Monitor vector growth: token_holdings and token_balances vectors could grow large, affecting account size limits

5. Add overflow tests: Create specific fuzz tests for edge cases around u64::MAX values

Test Scenarios for Fuzzing

1. Maximum value deposits: Test deposits of u64::MAX - 1
2. Cumulative overflow: Multiple maximum deposits to trigger overflow
3. Fee calculation edge cases: Investments near u64::MAX / 2
4. Token allocation precision: Large token amounts with small SOL amounts
5. Round counter stress test: Rapid round increments (though impractical)
6. Vector size limits: Adding maximum number of token types

SECURITY_AUDIT_REPORT.md

Smart Contract Security Audit Report

Sol-Vault Protocol - Critical Vulnerabilities Found

Audit Date: July 2025 Test Scenarios: 5,000+ contract execution flows

---

Executive Summary

This security audit identified critical mathematical vulnerabilities in the Sol-Vault protocol through contract execution simulation. The fuzzing discovered genuine bugs that could lead to token theft and protocol insolvency.

Critical Findings Overview

Severity	Count	Impact
LOW	2	Precision/Rounding Issues

---

Methodology

This audit used contract execution simulation throw fuzzing techniques, 1,000+ randomly generated scenarios (fuzzing) with various user deposit/investment combinations.

---

🟢 LOW SEVERITY

L-1: Precision Loss in Large Value Operations

Attack Vector: Economic Model Manipulation

Description: Minor discrepancies between calculated user contributions and actual vault spending, indicating fee calculation errors. The token over-allocation could seem negligible as it only accounts for 1.1216477e-9 % of the total tokens received by the vault. Evidence: The output of cargo fuzz run update_user_balances_batch_fuzz -- -max_total_time=30 reads:

🚨 VULNERABILITY DETECTED 🚨
=====================================
💰 FINANCIAL METRICS:
  Sum of User Deposits:     189.210525595 SOL (189210525595 lamports)
  Sum of User Investments:  14.196834895 SOL (14196834895 lamports)
  Total SOL Spent:          13.912898198 SOL (13912898198 lamports)
  Accumulated Fees:         0.283936697 SOL (283936697 lamports)

🪙 TOKEN ALLOCATION:
  Tokens Received by Vault: 382473034410
  Tokens Distributed:       382473034839
  Over-allocation:          429 tokens

L-2: Unlikely division by zero

Risk Level: LOW Attack Vector: Edge Case Exploitation

Description: The token allocation formula can divide by zero when total_sol_spent = 0 but users have contributions. The edge case is technically valid but extremely unlikely in practice - would require very specific small amounts that result in zero net spending after fees

When vault.next_round_buy_amount is very small (< 50 lamports):

vault.next_round_buy_amount = 49 lamports
fee_amount = (49 * 2) / 100 = 98 / 100 = 0 // (integer division truncates)
net_buy_amount = 49 - 0 = 49 lamports

However, if the fee calculation somehow results in net_buy_amount = 0, then current_round_spend = 0.

---

state.rs

use anchor_lang::prelude::*;
// use std::str::FromStr;

pub const VAULT_SEED: &[u8] = b"collective_vault";
pub const USER_SEED: &[u8] = b"user";

// Jupiter program ID TODO: use this instead of the one being passed by the client
 // pub fn jupiter_program_id() -> Pubkey {
//     Pubkey::from_str("JUP6LkbZbjS1jKKwapdHNy74zcZ3tLUZoi5QNyVTaV4").unwrap()
// }

#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
pub enum VaultStage {
    Deposit,   // Only deposits allowed
    Buy,       // Execute collective buys and sells
    Update,    // Update user balances in batches
    Withdraw,  // Users can withdraw
}

pub fn require_authority_role(caller: &Pubkey, vault: &Vault) -> Result<()> {
    require!(
        *caller == vault.authority,
        VaultErrorCode::UnauthorizedCaller
    );
    Ok(())
}

pub fn require_executor_role(caller: &Pubkey, vault: &Vault) -> Result<()> {
    require!(
        *caller == vault.authority || vault.executor.map_or(false, |exec| *caller == exec),
        VaultErrorCode::UnauthorizedCaller
    );
    Ok(())
}

pub fn require_stage(vault: &Vault, expected_stage: VaultStage) -> Result<()> {
    require!(
        vault.stage == expected_stage,
        VaultErrorCode::WrongStage
    );
    Ok(())
}

// Helper function to deserialize and then serialize back an account after modification
// This is useful when you need to modify an account from remaining_accounts
pub fn modify_account<T, F>(
    account_info: &AccountInfo,
    modifier: F,
) -> Result<()>
where
    T: AccountDeserialize + AccountSerialize,
    F: FnOnce(&mut T) -> Result<()>,
{
    // Step 1: Load account data from blockchain storage
    let mut account_data = account_info.try_borrow_mut_data()?;

    // Step 2: Convert raw bytes into Rust struct
    let mut account = T::try_deserialize(&mut account_data.as_ref())?;

    // Step 3: Apply the modification function
    modifier(&mut account)?;

    // Step 4: Convert modified struct back to bytes
    let mut updated_account_bytes = Vec::new();
    account.try_serialize(&mut updated_account_bytes)?;

    // Step 5: Ensure new data fits in existing account space
    require!(
        updated_account_bytes.len() <= account_data.len(),
        VaultErrorCode::InvalidAccountsProvided
    );

    // Step 6: Write updated bytes back to blockchain storage
    account_data[..updated_account_bytes.len()].copy_from_slice(&updated_account_bytes);

    Ok(())
}

#[derive(AnchorSerialize, AnchorDeserialize, Clone)]
pub struct RoundSnapshot {
    pub round: u64,
    pub total_sol_spent: u64,
    pub total_tokens_received: u64,
    pub token_mint: Pubkey,
}

#[account]
pub struct Vault {
    pub authority: Pubkey,
    pub executor: Option<Pubkey>,
    pub deposits: u64,
    pub next_round_buy_amount: u64,
    pub current_round_spend: u64,
    pub token_holdings: Vec<TokenBalance>,
    pub stage: VaultStage,
    pub current_round: u64,
    pub accumulated_fees: u64,
    // New field for O(1) balance updates
    pub round_snapshots: Vec<RoundSnapshot>,
}

impl Vault {

    pub fn new(authority: Pubkey) -> Self {
        Self {
            authority,
            executor: None,
            deposits: 0,
            next_round_buy_amount: 0,
            current_round_spend: 0,
            token_holdings: Vec::new(),
            stage: VaultStage::Deposit,
            current_round: 0,
            accumulated_fees: 0,
            round_snapshots: Vec::new(),
        }
    }

    pub fn add_deposits(&mut self, amount: u64) -> Result<()> {
        self.deposits = self.deposits
            .checked_add(amount)
            .ok_or(VaultErrorCode::MathOverflow)?;
        Ok(())
    }

    pub fn set_current_spend(&mut self, amount: u64) {
        self.current_round_spend = amount;
    }

    pub fn add_fees(&mut self, amount: u64) -> Result<()> {
        self.accumulated_fees = self.accumulated_fees
            .checked_add(amount)
            .ok_or(VaultErrorCode::MathOverflow)?;
        Ok(())
    }

    pub fn subtract_fees(&mut self, amount: u64) -> Result<()> {
        require!(
            self.accumulated_fees >= amount,
            VaultErrorCode::InsufficientFunds
        );
        self.accumulated_fees -= amount;
        Ok(())
    }

    pub fn add_to_next_round(&mut self, amount: u64) -> Result<()> {
        self.next_round_buy_amount = self.next_round_buy_amount
            .checked_add(amount)
            .ok_or(VaultErrorCode::MathOverflow)?;
        Ok(())
    }

    pub fn subtract_from_next_round(&mut self, amount: u64) -> Result<()> {
        self.next_round_buy_amount = self.next_round_buy_amount
            .checked_sub(amount)
            .ok_or(VaultErrorCode::MathOverflow)?;
        Ok(())
    }

    pub fn update_next_round_amount(&mut self, old_contribution: u64, new_contribution: u64) -> Result<()> {
        match new_contribution.cmp(&old_contribution) {
            std::cmp::Ordering::Greater => {
                let delta = new_contribution - old_contribution;
                self.add_to_next_round(delta)
            }
            std::cmp::Ordering::Less => {
                let delta = old_contribution - new_contribution;
                self.subtract_from_next_round(delta)
            }
            std::cmp::Ordering::Equal => Ok(()),
        }
    }

    // pub fn add_to_next_round(&mut self, amount: u64) -> Result<()> {
    //     self.next_round_buy_amount = self.next_round_buy_amount
    //         .checked_add(amount)
    //         .ok_or(VaultErrorCode::MathOverflow)?;
    //     Ok(())
    // }

    pub fn add_or_update_token_balance(&mut self, token_mint: Pubkey, amount: u64) -> Result<()> {
        // Find existing token balance
        for token_balance in self.token_holdings.iter_mut() {
            if token_balance.mint == token_mint {
                token_balance.amount = token_balance.amount
                    .checked_add(amount)
                    .ok_or(VaultErrorCode::MathOverflow)?;
                return Ok(());
            }
        }

        // Token not found, add new entry
        self.token_holdings.push(TokenBalance {
            mint: token_mint,
            amount,
        });

        Ok(())
    }

    pub fn get_token_balance(&self, token_mint: Pubkey) -> Result<u64> {
        for token_balance in self.token_holdings.iter() {
            if token_balance.mint == token_mint {
                return Ok(token_balance.amount);
            }
        }
        Ok(0) // Token not found, return 0
    }

    pub fn subtract_token_balance(&mut self, token_mint: Pubkey, amount: u64) -> Result<()> {
        for token_balance in self.token_holdings.iter_mut() {
            if token_balance.mint == token_mint {
                require!(
                    token_balance.amount >= amount,
                    VaultErrorCode::InsufficientTokenBalance
                );
                token_balance.amount -= amount;
                return Ok(());
            }
        }
        return Err(VaultErrorCode::TokenNotFound.into());
    }

    // pub fn get_buy_amount(&self) -> u64 {
    //     self.total_buy_amount
    // }

    // pub fn subtract_buy_amount(&mut self, amount: u64) -> Result<()> {
    //     self.total_buy_amount = self.total_buy_amount
    //         .checked_sub(amount)
    //         .ok_or(VaultErrorCode::MathOverflow)?;
    //     Ok(())
    // }
}

#[derive(AnchorSerialize, AnchorDeserialize, Clone)]
pub struct TokenBalance {
    pub mint: Pubkey,
    pub amount: u64,
}

// #[derive(AnchorSerialize, AnchorDeserialize, Clone)]
// pub struct UserData {
//     pub user: Pubkey,
//     pub sol_balance: u64, // SOL balance for deposits/withdrawals
//     pub token_balances: Vec<TokenBalance>, // All token holdings
//     pub amount_per_round: u64, // How much SOL user wants to invest per round
//     pub last_participated_round: u64, // Last round user participated in
//     pub is_active: bool, // Whether user is currently participating
//     pub total_rounds_participated: u32, // Total rounds user has participated
//     pub reserved_amount_current_round: u64, // SOL reserved for current buy
//     pub is_reserved: bool, // Whether user has funds reserved for current round
// }

// impl UserData {
//     pub fn add_or_update_token_balance(&mut self, token_mint: Pubkey, amount: u64) -> Result<()> {
//         // Find existing token balance
//         for token_balance in self.token_balances.iter_mut() {
//             if token_balance.mint == token_mint {
//                 token_balance.amount = token_balance.amount
//                     .checked_add(amount)
//                     .ok_or(VaultErrorCode::MathOverflow)?;
//                 return Ok(());
//             }
//         }

//         // Token not found, add new entry
//         self.token_balances.push(TokenBalance {
//             mint: token_mint,
//             amount,
//         });

//         Ok(())
//     }

//     pub fn get_token_balance(&self, token_mint: Pubkey) -> Result<u64> {
//         for token_balance in self.token_balances.iter() {
//             if token_balance.mint == token_mint {
//                 return Ok(token_balance.amount);
//             }
//         }
//         Ok(0) // Token not found, return 0
//     }

//     pub fn subtract_token_balance(&mut self, token_mint: Pubkey, amount: u64) -> Result<()> {
//         for token_balance in self.token_balances.iter_mut() {
//             if token_balance.mint == token_mint {
//                 require!(
//                     token_balance.amount >= amount,
//                     VaultErrorCode::InsufficientTokenBalance
//                 );
//                 token_balance.amount -= amount;
//                 return Ok(());
//             }
//         }
//         return Err(VaultErrorCode::TokenNotFound.into());
//     }
// }

#[account]
pub struct UserAccount {
    pub user: Pubkey,
    pub sol_balance: u64,
    pub token_balances: Vec<TokenBalance>,
    pub investment_per_token: u64,
    pub is_active: bool,
    // New field to track last processed round for lazy evaluation
    pub last_processed_round: u64,
}

impl UserAccount {

    pub fn new(user: Pubkey) -> Self {
        Self {
            user,
            sol_balance: 0,
            token_balances: Vec::new(),
            investment_per_token: 0,
            is_active: false,
            last_processed_round: 0,
        }
    }

    // ✅ O(1) Lazy Balance Calculation
    pub fn calculate_pending_tokens(&self, vault: &Vault, target_round: u64) -> Result<Vec<TokenBalance>> {
        let mut pending_tokens = Vec::new();

        // Process all rounds since last update
        for snapshot in vault.round_snapshots.iter() {
            if snapshot.round > self.last_processed_round && snapshot.round <= target_round {
                let contribution = if self.sol_balance >= self.investment_per_token && self.is_active {
                    self.investment_per_token
                } else {
                    0
                };

                if contribution > 0 {
                    let user_fee = (contribution * 2) / 100;
                    let user_net_contribution = contribution - user_fee;
                    let tokens = (user_net_contribution as u128 * snapshot.total_tokens_received as u128 / snapshot.total_sol_spent as u128) as u64;

                    // Add to pending tokens
                    let mut found = false;
                    for pending in pending_tokens.iter_mut() {
                        if pending.mint == snapshot.token_mint {
                            pending.amount += tokens;
                            found = true;
                            break;
                        }
                    }
                    if !found {
                        pending_tokens.push(TokenBalance {
                            mint: snapshot.token_mint,
                            amount: tokens,
                        });
                    }
                }
            }
        }

        Ok(pending_tokens)
    }

    // ✅ Apply pending balance updates when needed
    pub fn apply_pending_updates(&mut self, vault: &Vault, target_round: u64) -> Result<()> {
        let pending_tokens = self.calculate_pending_tokens(vault, target_round)?;

        // Apply SOL deductions and token additions
        for snapshot in vault.round_snapshots.iter() {
            if snapshot.round > self.last_processed_round && snapshot.round <= target_round {
                let contribution = if self.sol_balance >= self.investment_per_token && self.is_active {
                    self.investment_per_token
                } else {
                    0
                };

                if contribution > 0 {
                    self.subtract_sol_balance(contribution)?;
                }
            }
        }

        // Add tokens
        for pending in pending_tokens {
            self.add_or_update_token_balance(pending.mint, pending.amount)?;
        }

        self.last_processed_round = target_round;
        Ok(())
    }

    pub fn is_uninitialized(&self) -> bool {
        self.user == Pubkey::default()
    }

    pub fn add_or_update_token_balance(&mut self, token_mint: Pubkey, amount: u64) -> Result<()> {
        for token_balance in self.token_balances.iter_mut() {
            if token_balance.mint == token_mint {
                token_balance.amount = token_balance.amount
                    .checked_add(amount)
                    .ok_or(VaultErrorCode::MathOverflow)?;
                return Ok(());
            }
        }

        self.token_balances.push(TokenBalance {
            mint: token_mint,
            amount,
        });

        Ok(())
    }

    pub fn get_token_balance(&self, token_mint: Pubkey) -> Result<u64> {
        for token_balance in self.token_balances.iter() {
            if token_balance.mint == token_mint {
                return Ok(token_balance.amount);
            }
        }
        Ok(0)
    }

    pub fn subtract_token_balance(&mut self, token_mint: Pubkey, amount: u64) -> Result<()> {
        for token_balance in self.token_balances.iter_mut() {
            if token_balance.mint == token_mint {
                require!(
                    token_balance.amount >= amount,
                    VaultErrorCode::InsufficientTokenBalance
                );
                token_balance.amount -= amount;
                return Ok(());
            }
        }
        return Err(VaultErrorCode::TokenNotFound.into());
    }

    pub fn add_sol_balance(&mut self, amount: u64) -> Result<()> {
        self.sol_balance = self.sol_balance
            .checked_add(amount)
            .ok_or(VaultErrorCode::MathOverflow)?;
        Ok(())
    }

    pub fn subtract_sol_balance(&mut self, amount: u64) -> Result<()> {
        require!(
            self.sol_balance >= amount,
            VaultErrorCode::InsufficientBalance
        );
        self.sol_balance -= amount;
        Ok(())
    }

    pub fn get_sol_balance(&self) -> u64 {
        self.sol_balance
    }
}

#[derive(AnchorSerialize, AnchorDeserialize, Clone)]
pub struct UserPortfolio {
    pub user: Pubkey,
    pub sol_balance: u64,
    pub token_balances: Vec<TokenBalance>,
    pub amount_per_round: u64,
    pub is_active: bool,
    pub last_participated_round: u64,
    pub total_rounds_participated: u32,
}

#[derive(AnchorSerialize, AnchorDeserialize, Clone)]
pub struct VaultInfo {
    pub authority: Pubkey,
    pub executor: Option<Pubkey>,
    pub total_deposits: u64,
    pub token_holdings: Vec<TokenBalance>,
    pub current_round: u64,
}

// Data structures for instruction parameters
#[derive(AnchorSerialize, AnchorDeserialize, Clone)]
pub struct UserBuyData {
    pub user: Pubkey,
    pub amount_spent: u64,
    pub tokens_received: u64,
}

// #[account]
// pub struct UserBalance {
//     pub user: Pubkey, // TODO: evaluate remove this
//     pub balance: u64, // SOL balance
//     pub token_balance: u64, // Token balance
// }

// impl UserBalance {
//     pub fn initialize_if_needed(&mut self, user: Pubkey) -> bool {
//         let is_first_deposit = self.user == Pubkey::default();
//         if is_first_deposit {
//             *self = UserBalance {
//                 user,
//                 balance: 0,
//                 token_balance: 0,
//             };
//         }
//         is_first_deposit
//     }

//     pub fn add_balance(&mut self, amount: u64) -> Result<()> {
//         self.balance = self.balance
//             .checked_add(amount)
//             .ok_or(VaultErrorCode::MathOverflow)?;
//         Ok(())
//     }
// }

#[error_code]
pub enum VaultErrorCode {
    #[msg("Insufficient balance")]
    InsufficientBalance,
    #[msg("User not found")]
    UserNotFound,
    #[msg("Insufficient vault balance")]
    InsufficientVaultBalance,
    #[msg("Destination overflow")]
    DestinationOverflow,
    #[msg("Unauthorized to update user balance")]
    UnauthorizedBalanceUpdate,
    #[msg("User mismatch")]
    UserMismatch,
    #[msg("Insufficient token balance")]
    InsufficientTokenBalance,
    #[msg("Trading not active")]
    TradingNotActive,
    #[msg("Invalid Jupiter program")]
    InvalidJupiterProgram,
    #[msg("Insufficient funds for operation")]
    InsufficientFunds,
    #[msg("Invalid calculation")]
    InvalidCalculation,
    #[msg("Insufficient tokens received")]
    InsufficientTokensReceived,
    #[msg("Invalid destination")]
    InvalidDestination,
    #[msg("Caller is not the vault authority or executor")]
    UnauthorizedCaller,
    #[msg("Math operation overflow")]
    MathOverflow,
    #[msg("Invalid accounts provided")]
    InvalidAccountsProvided,
    #[msg("Invalid account type")]
    InvalidAccountType,
    #[msg("User not eligible for this round")]
    UserNotEligibleForRound,
    #[msg("Reservation amount mismatch")]
    ReservationMismatch,
    #[msg("No reserved funds found")]
    NoReservedFunds,
    #[msg("Token not found")]
    TokenNotFound,
    #[msg("Wrong stage for this operation")]
    WrongStage,
    #[msg("Setup not completed")]
    SetupNotCompleted,
    #[msg("No users in batch")]
    EmptyBatch,
}