Staking Program

┌─────────────────────────────────────────────────────────────────┐
│                     Staking Flow                                │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  Pool Admin                                                     │
│     │                                                           │
│     ▼                                                           │
│  ┌─────────────────────┐                                        │
│  │   Staking Pool      │◄──── Reward Token Deposits             │
│  │   ────────────────  │                                        │
│  │   Total Staked      │                                        │
│  │   Reward Rate       │                                        │
│  │   Last Update Time  │                                        │
│  └─────────────────────┘                                        │
│           │                                                     │
│           ▼                                                     │
│  ┌─────────────────────┐     ┌─────────────────────┐            │
│  │   User Stake #1     │     │   User Stake #2     │            │
│  │   ────────────────  │     │   ────────────────  │            │
│  │   Amount: 1000      │     │   Amount: 500       │            │
│  │   Rewards: 50       │     │   Rewards: 25       │            │
│  └─────────────────────┘     └─────────────────────┘            │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

use anchor_lang::prelude::*;
use anchor_spl::token::{self, Token, TokenAccount, Transfer, Mint};
use anchor_spl::associated_token::AssociatedToken;

declare_id!("Stake11111111111111111111111111111111111111");

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

    /// Initialize a new staking pool
    pub fn initialize_pool(
        ctx: Context<InitializePool>,
        reward_rate: u64,  // Rewards per second per staked token (scaled by 1e9)
        lock_period: i64,  // Minimum stake duration in seconds
        min_stake: u64,    // Minimum stake amount
    ) -> Result<()> {
        let pool = &mut ctx.accounts.pool;
        pool.authority = ctx.accounts.authority.key();
        pool.stake_mint = ctx.accounts.stake_mint.key();
        pool.reward_mint = ctx.accounts.reward_mint.key();
        pool.reward_rate = reward_rate;
        pool.lock_period = lock_period;
        pool.min_stake = min_stake;
        pool.total_staked = 0;
        pool.reward_per_token_stored = 0;
        pool.last_update_time = Clock::get()?.unix_timestamp;
        pool.paused = false;
        pool.bump = ctx.bumps.pool;
        pool.stake_vault_bump = ctx.bumps.stake_vault;
        pool.reward_vault_bump = ctx.bumps.reward_vault;

        emit!(PoolInitialized {
            pool: ctx.accounts.pool.key(),
            stake_mint: pool.stake_mint,
            reward_mint: pool.reward_mint,
            reward_rate,
        });

        Ok(())
    }

    /// Stake tokens
    pub fn stake(ctx: Context<Stake>, amount: u64) -> Result<()> {
        let pool = &ctx.accounts.pool;
        require!(!pool.paused, StakingError::PoolPaused);
        require!(amount >= pool.min_stake, StakingError::BelowMinimumStake);

        // Update rewards before modifying stake
        update_rewards(&mut ctx.accounts.pool, Some(&mut ctx.accounts.user_stake))?;

        // Transfer stake tokens to vault
        let cpi_accounts = Transfer {
            from: ctx.accounts.user_stake_token.to_account_info(),
            to: ctx.accounts.stake_vault.to_account_info(),
            authority: ctx.accounts.user.to_account_info(),
        };
        let cpi_ctx = CpiContext::new(
            ctx.accounts.token_program.to_account_info(),
            cpi_accounts,
        );
        token::transfer(cpi_ctx, amount)?;

        // Update user stake
        let user_stake = &mut ctx.accounts.user_stake;
        user_stake.staked_amount = user_stake.staked_amount.checked_add(amount)
            .ok_or(StakingError::Overflow)?;
        user_stake.last_stake_time = Clock::get()?.unix_timestamp;

        // Update pool
        let pool = &mut ctx.accounts.pool;
        pool.total_staked = pool.total_staked.checked_add(amount)
            .ok_or(StakingError::Overflow)?;

        emit!(Staked {
            pool: ctx.accounts.pool.key(),
            user: ctx.accounts.user.key(),
            amount,
            total_staked: user_stake.staked_amount,
        });

        Ok(())
    }

    /// Unstake tokens
    pub fn unstake(ctx: Context<Unstake>, amount: u64) -> Result<()> {
        let user_stake = &ctx.accounts.user_stake;
        let pool = &ctx.accounts.pool;

        require!(amount > 0, StakingError::InvalidAmount);
        require!(amount <= user_stake.staked_amount, StakingError::InsufficientStake);

        // Check lock period
        let clock = Clock::get()?;
        let time_since_stake = clock.unix_timestamp - user_stake.last_stake_time;
        require!(
            time_since_stake >= pool.lock_period,
            StakingError::StakeLocked
        );

        // Update rewards
        update_rewards(&mut ctx.accounts.pool, Some(&mut ctx.accounts.user_stake))?;

        // Transfer tokens back to user
        let pool_key = ctx.accounts.pool.key();
        let seeds = &[
            b"pool",
            pool.stake_mint.as_ref(),
            pool.reward_mint.as_ref(),
            &[pool.bump],
        ];
        let signer_seeds = &[&seeds[..]];

        let cpi_accounts = Transfer {
            from: ctx.accounts.stake_vault.to_account_info(),
            to: ctx.accounts.user_stake_token.to_account_info(),
            authority: ctx.accounts.pool.to_account_info(),
        };
        let cpi_ctx = CpiContext::new_with_signer(
            ctx.accounts.token_program.to_account_info(),
            cpi_accounts,
            signer_seeds,
        );
        token::transfer(cpi_ctx, amount)?;

        // Update user stake
        let user_stake = &mut ctx.accounts.user_stake;
        user_stake.staked_amount = user_stake.staked_amount.checked_sub(amount)
            .ok_or(StakingError::Overflow)?;

        // Update pool
        let pool = &mut ctx.accounts.pool;
        pool.total_staked = pool.total_staked.checked_sub(amount)
            .ok_or(StakingError::Overflow)?;

        emit!(Unstaked {
            pool: ctx.accounts.pool.key(),
            user: ctx.accounts.user.key(),
            amount,
            remaining: user_stake.staked_amount,
        });

        Ok(())
    }

    /// Claim accumulated rewards
    pub fn claim_rewards(ctx: Context<ClaimRewards>) -> Result<()> {
        // Update rewards
        update_rewards(&mut ctx.accounts.pool, Some(&mut ctx.accounts.user_stake))?;

        let rewards = ctx.accounts.user_stake.pending_rewards;
        require!(rewards > 0, StakingError::NoRewardsToClaim);

        // Check reward vault has enough tokens
        require!(
            ctx.accounts.reward_vault.amount >= rewards,
            StakingError::InsufficientRewardBalance
        );

        // Transfer rewards
        let pool = &ctx.accounts.pool;
        let seeds = &[
            b"pool",
            pool.stake_mint.as_ref(),
            pool.reward_mint.as_ref(),
            &[pool.bump],
        ];
        let signer_seeds = &[&seeds[..]];

        let cpi_accounts = Transfer {
            from: ctx.accounts.reward_vault.to_account_info(),
            to: ctx.accounts.user_reward_token.to_account_info(),
            authority: ctx.accounts.pool.to_account_info(),
        };
        let cpi_ctx = CpiContext::new_with_signer(
            ctx.accounts.token_program.to_account_info(),
            cpi_accounts,
            signer_seeds,
        );
        token::transfer(cpi_ctx, rewards)?;

        // Reset pending rewards
        let user_stake = &mut ctx.accounts.user_stake;
        user_stake.pending_rewards = 0;

        emit!(RewardsClaimed {
            pool: ctx.accounts.pool.key(),
            user: ctx.accounts.user.key(),
            amount: rewards,
        });

        Ok(())
    }

    /// Compound rewards into stake (for same-token staking)
    pub fn compound(ctx: Context<Compound>) -> Result<()> {
        let pool = &ctx.accounts.pool;
        require!(
            pool.stake_mint == pool.reward_mint,
            StakingError::CannotCompound
        );
        require!(!pool.paused, StakingError::PoolPaused);

        // Update rewards
        update_rewards(&mut ctx.accounts.pool, Some(&mut ctx.accounts.user_stake))?;

        let rewards = ctx.accounts.user_stake.pending_rewards;
        require!(rewards > 0, StakingError::NoRewardsToClaim);

        // Transfer from reward vault to stake vault
        let pool = &ctx.accounts.pool;
        let seeds = &[
            b"pool",
            pool.stake_mint.as_ref(),
            pool.reward_mint.as_ref(),
            &[pool.bump],
        ];
        let signer_seeds = &[&seeds[..]];

        let cpi_accounts = Transfer {
            from: ctx.accounts.reward_vault.to_account_info(),
            to: ctx.accounts.stake_vault.to_account_info(),
            authority: ctx.accounts.pool.to_account_info(),
        };
        let cpi_ctx = CpiContext::new_with_signer(
            ctx.accounts.token_program.to_account_info(),
            cpi_accounts,
            signer_seeds,
        );
        token::transfer(cpi_ctx, rewards)?;

        // Update balances
        let user_stake = &mut ctx.accounts.user_stake;
        user_stake.staked_amount = user_stake.staked_amount.checked_add(rewards)
            .ok_or(StakingError::Overflow)?;
        user_stake.pending_rewards = 0;

        let pool = &mut ctx.accounts.pool;
        pool.total_staked = pool.total_staked.checked_add(rewards)
            .ok_or(StakingError::Overflow)?;

        emit!(Compounded {
            pool: ctx.accounts.pool.key(),
            user: ctx.accounts.user.key(),
            amount: rewards,
            new_stake: user_stake.staked_amount,
        });

        Ok(())
    }

    /// Admin: Add rewards to pool
    pub fn fund_rewards(ctx: Context<FundRewards>, amount: u64) -> Result<()> {
        let cpi_accounts = Transfer {
            from: ctx.accounts.funder_token.to_account_info(),
            to: ctx.accounts.reward_vault.to_account_info(),
            authority: ctx.accounts.funder.to_account_info(),
        };
        let cpi_ctx = CpiContext::new(
            ctx.accounts.token_program.to_account_info(),
            cpi_accounts,
        );
        token::transfer(cpi_ctx, amount)?;

        emit!(RewardsFunded {
            pool: ctx.accounts.pool.key(),
            funder: ctx.accounts.funder.key(),
            amount,
        });

        Ok(())
    }

    /// Admin: Update pool parameters
    pub fn update_pool(
        ctx: Context<UpdatePool>,
        new_reward_rate: Option<u64>,
        new_lock_period: Option<i64>,
        new_min_stake: Option<u64>,
    ) -> Result<()> {
        // Update rewards with current rate first
        update_rewards(&mut ctx.accounts.pool, None)?;

        let pool = &mut ctx.accounts.pool;

        if let Some(rate) = new_reward_rate {
            pool.reward_rate = rate;
        }
        if let Some(period) = new_lock_period {
            pool.lock_period = period;
        }
        if let Some(min) = new_min_stake {
            pool.min_stake = min;
        }

        emit!(PoolUpdated {
            pool: ctx.accounts.pool.key(),
            reward_rate: pool.reward_rate,
            lock_period: pool.lock_period,
            min_stake: pool.min_stake,
        });

        Ok(())
    }

    /// Admin: Pause/unpause pool
    pub fn set_paused(ctx: Context<SetPaused>, paused: bool) -> Result<()> {
        ctx.accounts.pool.paused = paused;
        Ok(())
    }
}

// Helper functions

fn update_rewards(pool: &mut Account<Pool>, user_stake: Option<&mut Account<UserStake>>) -> Result<()> {
    let clock = Clock::get()?;
    let current_time = clock.unix_timestamp;

    if pool.total_staked > 0 {
        let time_elapsed = current_time - pool.last_update_time;
        let new_rewards = (time_elapsed as u128)
            .checked_mul(pool.reward_rate as u128)
            .ok_or(StakingError::Overflow)?;

        let reward_per_token_delta = new_rewards
            .checked_mul(PRECISION as u128)
            .ok_or(StakingError::Overflow)?
            .checked_div(pool.total_staked as u128)
            .unwrap_or(0);

        pool.reward_per_token_stored = pool.reward_per_token_stored
            .checked_add(reward_per_token_delta as u64)
            .ok_or(StakingError::Overflow)?;
    }

    pool.last_update_time = current_time;

    if let Some(user) = user_stake {
        let earned = calculate_earned(pool, user)?;
        user.pending_rewards = user.pending_rewards
            .checked_add(earned)
            .ok_or(StakingError::Overflow)?;
        user.reward_per_token_paid = pool.reward_per_token_stored;
    }

    Ok(())
}

fn calculate_earned(pool: &Account<Pool>, user: &Account<UserStake>) -> Result<u64> {
    let reward_delta = pool.reward_per_token_stored
        .checked_sub(user.reward_per_token_paid)
        .ok_or(StakingError::Overflow)?;

    let earned = (user.staked_amount as u128)
        .checked_mul(reward_delta as u128)
        .ok_or(StakingError::Overflow)?
        .checked_div(PRECISION as u128)
        .unwrap_or(0) as u64;

    Ok(earned)
}

// Constants
const PRECISION: u64 = 1_000_000_000; // 1e9 for reward calculations

// Accounts

#[derive(Accounts)]
pub struct InitializePool<'info> {
    #[account(
        init,
        payer = authority,
        space = 8 + Pool::INIT_SPACE,
        seeds = [b"pool", stake_mint.key().as_ref(), reward_mint.key().as_ref()],
        bump
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        init,
        payer = authority,
        token::mint = stake_mint,
        token::authority = pool,
        seeds = [b"stake_vault", pool.key().as_ref()],
        bump
    )]
    pub stake_vault: Account<'info, TokenAccount>,

    #[account(
        init,
        payer = authority,
        token::mint = reward_mint,
        token::authority = pool,
        seeds = [b"reward_vault", pool.key().as_ref()],
        bump
    )]
    pub reward_vault: Account<'info, TokenAccount>,

    pub stake_mint: Account<'info, Mint>,
    pub reward_mint: Account<'info, Mint>,

    #[account(mut)]
    pub authority: Signer<'info>,

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

#[derive(Accounts)]
pub struct Stake<'info> {
    #[account(
        mut,
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        init_if_needed,
        payer = user,
        space = 8 + UserStake::INIT_SPACE,
        seeds = [b"user_stake", pool.key().as_ref(), user.key().as_ref()],
        bump
    )]
    pub user_stake: Account<'info, UserStake>,

    #[account(
        mut,
        seeds = [b"stake_vault", pool.key().as_ref()],
        bump = pool.stake_vault_bump,
    )]
    pub stake_vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        token::mint = pool.stake_mint,
        token::authority = user,
    )]
    pub user_stake_token: Account<'info, TokenAccount>,

    #[account(mut)]
    pub user: Signer<'info>,

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

#[derive(Accounts)]
pub struct Unstake<'info> {
    #[account(
        mut,
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [b"user_stake", pool.key().as_ref(), user.key().as_ref()],
        bump,
        has_one = pool,
    )]
    pub user_stake: Account<'info, UserStake>,

    #[account(
        mut,
        seeds = [b"stake_vault", pool.key().as_ref()],
        bump = pool.stake_vault_bump,
    )]
    pub stake_vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        token::mint = pool.stake_mint,
        token::authority = user,
    )]
    pub user_stake_token: Account<'info, TokenAccount>,

    #[account(mut)]
    pub user: Signer<'info>,

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

#[derive(Accounts)]
pub struct ClaimRewards<'info> {
    #[account(
        mut,
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [b"user_stake", pool.key().as_ref(), user.key().as_ref()],
        bump,
        has_one = pool,
    )]
    pub user_stake: Account<'info, UserStake>,

    #[account(
        mut,
        seeds = [b"reward_vault", pool.key().as_ref()],
        bump = pool.reward_vault_bump,
    )]
    pub reward_vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        token::mint = pool.reward_mint,
    )]
    pub user_reward_token: Account<'info, TokenAccount>,

    #[account(mut)]
    pub user: Signer<'info>,

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

#[derive(Accounts)]
pub struct Compound<'info> {
    #[account(
        mut,
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [b"user_stake", pool.key().as_ref(), user.key().as_ref()],
        bump,
        has_one = pool,
    )]
    pub user_stake: Account<'info, UserStake>,

    #[account(
        mut,
        seeds = [b"stake_vault", pool.key().as_ref()],
        bump = pool.stake_vault_bump,
    )]
    pub stake_vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        seeds = [b"reward_vault", pool.key().as_ref()],
        bump = pool.reward_vault_bump,
    )]
    pub reward_vault: Account<'info, TokenAccount>,

    #[account(mut)]
    pub user: Signer<'info>,

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

#[derive(Accounts)]
pub struct FundRewards<'info> {
    #[account(
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
    )]
    pub pool: Account<'info, Pool>,

    #[account(
        mut,
        seeds = [b"reward_vault", pool.key().as_ref()],
        bump = pool.reward_vault_bump,
    )]
    pub reward_vault: Account<'info, TokenAccount>,

    #[account(
        mut,
        token::mint = pool.reward_mint,
        token::authority = funder,
    )]
    pub funder_token: Account<'info, TokenAccount>,

    #[account(mut)]
    pub funder: Signer<'info>,

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

#[derive(Accounts)]
pub struct UpdatePool<'info> {
    #[account(
        mut,
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
        has_one = authority,
    )]
    pub pool: Account<'info, Pool>,

    pub authority: Signer<'info>,
}

#[derive(Accounts)]
pub struct SetPaused<'info> {
    #[account(
        mut,
        seeds = [b"pool", pool.stake_mint.as_ref(), pool.reward_mint.as_ref()],
        bump = pool.bump,
        has_one = authority,
    )]
    pub pool: Account<'info, Pool>,

    pub authority: Signer<'info>,
}

// State

#[account]
#[derive(InitSpace)]
pub struct Pool {
    pub authority: Pubkey,
    pub stake_mint: Pubkey,
    pub reward_mint: Pubkey,
    pub reward_rate: u64,
    pub lock_period: i64,
    pub min_stake: u64,
    pub total_staked: u64,
    pub reward_per_token_stored: u64,
    pub last_update_time: i64,
    pub paused: bool,
    pub bump: u8,
    pub stake_vault_bump: u8,
    pub reward_vault_bump: u8,
}

#[account]
#[derive(InitSpace)]
pub struct UserStake {
    pub pool: Pubkey,
    pub owner: Pubkey,
    pub staked_amount: u64,
    pub pending_rewards: u64,
    pub reward_per_token_paid: u64,
    pub last_stake_time: i64,
}

// Events

#[event]
pub struct PoolInitialized {
    pub pool: Pubkey,
    pub stake_mint: Pubkey,
    pub reward_mint: Pubkey,
    pub reward_rate: u64,
}

#[event]
pub struct Staked {
    pub pool: Pubkey,
    pub user: Pubkey,
    pub amount: u64,
    pub total_staked: u64,
}

#[event]
pub struct Unstaked {
    pub pool: Pubkey,
    pub user: Pubkey,
    pub amount: u64,
    pub remaining: u64,
}

#[event]
pub struct RewardsClaimed {
    pub pool: Pubkey,
    pub user: Pubkey,
    pub amount: u64,
}

#[event]
pub struct Compounded {
    pub pool: Pubkey,
    pub user: Pubkey,
    pub amount: u64,
    pub new_stake: u64,
}

#[event]
pub struct RewardsFunded {
    pub pool: Pubkey,
    pub funder: Pubkey,
    pub amount: u64,
}

#[event]
pub struct PoolUpdated {
    pub pool: Pubkey,
    pub reward_rate: u64,
    pub lock_period: i64,
    pub min_stake: u64,
}

// Errors

#[error_code]
pub enum StakingError {
    #[msg("Pool is paused")]
    PoolPaused,

    #[msg("Amount below minimum stake")]
    BelowMinimumStake,

    #[msg("Invalid amount")]
    InvalidAmount,

    #[msg("Insufficient stake balance")]
    InsufficientStake,

    #[msg("Stake is still locked")]
    StakeLocked,

    #[msg("No rewards to claim")]
    NoRewardsToClaim,

    #[msg("Insufficient reward balance in pool")]
    InsufficientRewardBalance,

    #[msg("Cannot compound: stake and reward tokens must match")]
    CannotCompound,

    #[msg("Arithmetic overflow")]
    Overflow,
}

import * as anchor from "@coral-xyz/anchor";
import { Program } from "@coral-xyz/anchor";
import { Staking } from "../target/types/staking";
import * as token from "@solana/spl-token";

export class StakingClient {
  constructor(
    private program: Program<Staking>,
    private provider: anchor.AnchorProvider,
  ) {}

  async initializePool(
    stakeMint: anchor.web3.PublicKey,
    rewardMint: anchor.web3.PublicKey,
    rewardRate: anchor.BN,
    lockPeriod: anchor.BN,
    minStake: anchor.BN,
  ): Promise<anchor.web3.PublicKey> {
    const [poolPda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("pool"), stakeMint.toBuffer(), rewardMint.toBuffer()],
      this.program.programId,
    );

    const [stakeVaultPda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("stake_vault"), poolPda.toBuffer()],
      this.program.programId,
    );

    const [rewardVaultPda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("reward_vault"), poolPda.toBuffer()],
      this.program.programId,
    );

    await this.program.methods
      .initializePool(rewardRate, lockPeriod, minStake)
      .accounts({
        pool: poolPda,
        stakeVault: stakeVaultPda,
        rewardVault: rewardVaultPda,
        stakeMint,
        rewardMint,
        authority: this.provider.wallet.publicKey,
        tokenProgram: token.TOKEN_PROGRAM_ID,
        systemProgram: anchor.web3.SystemProgram.programId,
        rent: anchor.web3.SYSVAR_RENT_PUBKEY,
      })
      .rpc();

    return poolPda;
  }

  async stake(pool: anchor.web3.PublicKey, amount: anchor.BN): Promise<void> {
    const poolData = await this.program.account.pool.fetch(pool);
    const user = this.provider.wallet.publicKey;

    const [userStakePda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("user_stake"), pool.toBuffer(), user.toBuffer()],
      this.program.programId,
    );

    const [stakeVaultPda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("stake_vault"), pool.toBuffer()],
      this.program.programId,
    );

    const userStakeToken = await token.getAssociatedTokenAddress(
      poolData.stakeMint,
      user,
    );

    await this.program.methods
      .stake(amount)
      .accounts({
        pool,
        userStake: userStakePda,
        stakeVault: stakeVaultPda,
        userStakeToken,
        user,
        tokenProgram: token.TOKEN_PROGRAM_ID,
        systemProgram: anchor.web3.SystemProgram.programId,
      })
      .rpc();
  }

  async unstake(pool: anchor.web3.PublicKey, amount: anchor.BN): Promise<void> {
    const poolData = await this.program.account.pool.fetch(pool);
    const user = this.provider.wallet.publicKey;

    const [userStakePda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("user_stake"), pool.toBuffer(), user.toBuffer()],
      this.program.programId,
    );

    const [stakeVaultPda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("stake_vault"), pool.toBuffer()],
      this.program.programId,
    );

    const userStakeToken = await token.getAssociatedTokenAddress(
      poolData.stakeMint,
      user,
    );

    await this.program.methods
      .unstake(amount)
      .accounts({
        pool,
        userStake: userStakePda,
        stakeVault: stakeVaultPda,
        userStakeToken,
        user,
        tokenProgram: token.TOKEN_PROGRAM_ID,
      })
      .rpc();
  }

  async claimRewards(pool: anchor.web3.PublicKey): Promise<void> {
    const poolData = await this.program.account.pool.fetch(pool);
    const user = this.provider.wallet.publicKey;

    const [userStakePda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("user_stake"), pool.toBuffer(), user.toBuffer()],
      this.program.programId,
    );

    const [rewardVaultPda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("reward_vault"), pool.toBuffer()],
      this.program.programId,
    );

    const userRewardToken = await token.getOrCreateAssociatedTokenAccount(
      this.provider.connection,
      (this.provider.wallet as anchor.Wallet).payer,
      poolData.rewardMint,
      user,
    );

    await this.program.methods
      .claimRewards()
      .accounts({
        pool,
        userStake: userStakePda,
        rewardVault: rewardVaultPda,
        userRewardToken: userRewardToken.address,
        user,
        tokenProgram: token.TOKEN_PROGRAM_ID,
      })
      .rpc();
  }

  async getPendingRewards(
    pool: anchor.web3.PublicKey,
    user: anchor.web3.PublicKey,
  ): Promise<anchor.BN> {
    const [userStakePda] = anchor.web3.PublicKey.findProgramAddressSync(
      [Buffer.from("user_stake"), pool.toBuffer(), user.toBuffer()],
      this.program.programId,
    );

    const poolData = await this.program.account.pool.fetch(pool);
    const userStake = await this.program.account.userStake.fetch(userStakePda);

    // Calculate pending rewards
    const currentTime = Math.floor(Date.now() / 1000);
    const timeElapsed = currentTime - poolData.lastUpdateTime.toNumber();

    const PRECISION = new anchor.BN(1_000_000_000);

    let rewardPerTokenStored = poolData.rewardPerTokenStored;
    if (poolData.totalStaked.gt(new anchor.BN(0))) {
      const newRewards = new anchor.BN(timeElapsed).mul(poolData.rewardRate);
      const rewardPerTokenDelta = newRewards
        .mul(PRECISION)
        .div(poolData.totalStaked);
      rewardPerTokenStored = rewardPerTokenStored.add(rewardPerTokenDelta);
    }

    const rewardDelta = rewardPerTokenStored.sub(userStake.rewardPerTokenPaid);
    const earned = userStake.stakedAmount.mul(rewardDelta).div(PRECISION);

    return userStake.pendingRewards.add(earned);
  }
}

┌─────────────────────────────────────────────────────────────────┐
│                  Reward Calculation                             │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  Accumulated Reward Per Token:                                  │
│                                                                 │
│                     reward_rate × time_elapsed                  │
│  reward_per_token = ─────────────────────────── × PRECISION     │
│                          total_staked                           │
│                                                                 │
│  User's Pending Rewards:                                        │
│                                                                 │
│  earned = user_stake × (reward_per_token - user_paid)           │
│           ─────────────────────────────────────────────         │
│                          PRECISION                              │
│                                                                 │
│  Where PRECISION = 10^9 for avoiding floating point             │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

function calculateAPY(
  rewardRate: number, // Rewards per second per token
  stakeMintDecimals: number,
  rewardMintDecimals: number,
): number {
  const SECONDS_PER_YEAR = 365.25 * 24 * 60 * 60;
  const PRECISION = 1e9;

  // Rewards per year per token (normalized)
  const rewardsPerYear = (rewardRate / PRECISION) * SECONDS_PER_YEAR;

  // Adjust for decimal differences
  const decimalAdjustment = Math.pow(
    10,
    stakeMintDecimals - rewardMintDecimals,
  );

  // APY as percentage
  return rewardsPerYear * decimalAdjustment * 100;
}