usdc-plus-exchange 0.1.3

use anchor_lang::prelude::*;
use crate::errors::ReflectErrorCodes;
use crate::reflect::{AutoCompound, deserialise_autocompound};
use crate::spl_mint::get_mint_supply;
use crate::drift::get_usdc_amount_drift_data;
use crate::ids;

// deposited_vault_value = What the vault OWES to users (liability side)
// last_pool_value = What's ACTUALLY in Drift after extraction (asset side)

/// Computes how much USDC would be returned for a given amount of receipt tokens
/// Formula: USDC = token_amount * deposited_vault_value / effective_supply
pub fn compute_usdc_from_tokens(
    token_amount: u64,
    deposited_vault_value: u64,
    effective_supply: u64,
) -> Result<u64> {
    if effective_supply == 0 {
        return Err(ReflectErrorCodes::MathError.into());
    }
    
    // Use u128 to prevent overflow during multiplication.
    let numerator = (token_amount as u128) * (deposited_vault_value as u128);
    let result = numerator / (effective_supply as u128);
    
    // Check if result fits in u64.
    if result > u64::MAX as u128 {
        return Err(ReflectErrorCodes::MathError.into());
    }    
    Ok(result as u64)
}

/// Computes how many receipt tokens would be minted for a given USDC deposit amount
/// Formula: tokens = usdc_amount * effective_supply / deposited_vault_value
pub fn compute_tokens_from_usdc(
    usdc_amount: u64,
    deposited_vault_value: u64,
    effective_supply: u64,
) -> Result<u64> {
    // First deposit gets 1:1 ratio.
    if effective_supply == 0 || deposited_vault_value == 0 {
        return Ok(usdc_amount);
    }
    
    // Existing formula for subsequent deposits.
    let numerator = (usdc_amount as u128) * (effective_supply as u128);
    let result = numerator / (deposited_vault_value as u128);
    
    if result > u64::MAX as u128 {
        return Err(ReflectErrorCodes::MathError.into());
    }
    
    Ok(result as u64)
}

/// Calculates the amount of USDC+ receipt tokens a user will receive for depositing USDC.
/// 
/// This function provides the real-time exchange rate by incorporating the latest yield
/// from Drift protocol before calculating the token amount.
/// 
/// ## Parameters
/// - `data_usdc_controller`: Raw bytes from the USDC strategy controller account.
/// - `data_spot_market_usdc`: Raw bytes from Drift's USDC spot market account containing
///   current market state and interest rates.
/// - `data_user_account`: Raw bytes from Drift user account that holds the protocol's USDC
///   deposits and accumulated yield.
/// - `usdc_amount`: Amount of USDC to deposit, scaled by 10^6 (1 USDC = 1_000_000)
/// 
/// ## Returns
/// Amount of USDC+ tokens the user will receive (with 6 decimals places)
/// 
/// ## Process
/// 1. Extract from accounts following data:
///     - AutoCompound state from controller account
///     - Drift usdc spot market state
///     - Reflect's Drift user account state
///     - USDC+ supply
/// 2. Calculate amount of USDC at Drift
/// 3. Updates the AutoCompound state with latest Drift usdc holdings
/// 4. Calculates USDC+ tokens based on the updated exchange rate
/// 
/// ## Formula
/// tokens = usdc_amount * effective_supply / deposited_vault_value
/// 
/// ## Example
/// 
/// // Deposit 100 USDC
/// let tokens = exchange_rate_usdc_input(
///     &controller_data,
///     &spot_market_data,
///     &user_account_data,
///     69_000_000  // 69 USDC with 6 decimals
/// )?;
/// 
pub fn exchange_rate_usdc_input(    
    data_usdc_controller: &[u8],
    data_spot_market_usdc: &[u8],
    data_user_account: &[u8],
    data_usdc_plus_mint: &[u8],
    usdc_amount: u64,     
) -> Result<u64> {    

    // Get autocompound and recipients.
    let mut auto_compound: AutoCompound = deserialise_autocompound(data_usdc_controller)?;

    // Get supply of USDC+ tokens.
    let usdc_plus_supply = get_mint_supply(data_usdc_plus_mint)?;

    // Get usdc amount in drift
    let usdc_amount_drift = get_usdc_amount_drift_data(data_user_account, data_spot_market_usdc)?;
    
    // Update pool value to latest value from drift.
    auto_compound.update_pool(usdc_amount_drift, &vec![10_000], usdc_plus_supply)?;

    let deposited_vault_value: u64 = auto_compound.deposited_vault_value;
    
    compute_tokens_from_usdc(usdc_amount, deposited_vault_value, usdc_plus_supply)
}


/// Calculates USDC+ tokens for a USDC deposit using Anchor account validation.
/// 
/// This wrapper function provides account validation before calculating the exchange rate,
/// ensuring the correct Reflect protocol accounts are being used.
/// 
/// ## Parameters
/// - `usdc_controller_account`: The USDC strategy controller AccountInfo, must match `ids::usdc_controller::ID`
/// - `spot_market_usdc_account`: Drift's USDC spot market AccountInfo, must match `ids::usdc_spot_market::ID`
/// - `reflect_user_account`: Reflect's Drift user AccountInfo, must match `ids::reflect_user_account_strategy_0::ID`
/// - `usdc_amount`: Amount of USDC to deposit, scaled by 10^6 (1 USDC = 1_000_000)
/// 
/// ## Returns
/// Amount of USDC+ tokens the user will receive (with 6 decimal places)
/// 
/// ## Security
/// - Validates all account addresses match expected protocol accounts
/// - Prevents account substitution attacks
/// - Returns `InvalidAccount` error if any account doesn't match
/// 
/// ## Example
/// ```ignore
/// let tokens = exchange_rate_usdc_input_accounts(
///     &ctx.accounts.usdc_controller,
///     &ctx.accounts.spot_market,
///     &ctx.accounts.drift_user,
///     69_000_000  // 69 USDC
/// )?;
/// ```
pub fn exchange_rate_usdc_accounts(
    usdc_controller_account: &AccountInfo,
    spot_market_usdc_account: &AccountInfo,
    reflect_user_account: &AccountInfo,
    usdc_plus_mint_account: &AccountInfo,
    usdc_amount: u64,     
) -> Result<u64> { 
    require!(usdc_controller_account.key == &ids::usdc_controller::ID, ReflectErrorCodes::InvalidUsdcControllerAccount);
    require!(spot_market_usdc_account.key == &ids::usdc_spot_market::ID, ReflectErrorCodes::InvalidUsdcSpotMarketAccount);
    require!(reflect_user_account.key == &ids::reflect_user_account_strategy_0::ID, ReflectErrorCodes::InvalidReflectUserAccount);
    exchange_rate_usdc_input(    
        &usdc_controller_account.data.borrow(),
        &spot_market_usdc_account.data.borrow(),
        &reflect_user_account.data.borrow(),
        &usdc_plus_mint_account.data.borrow(),
        usdc_amount,     
    )
}


/// Calculates the amount of USDC a user will receive when redeeming USDC+ receipt tokens.
/// 
/// This function provides the real-time redemption value by incorporating the latest yield
/// from Drift protocol before calculating the USDC amount.
/// 
/// ## Parameters
/// - `data_usdc_controller`: Raw bytes from the USDC strategy controller account.
/// - `data_spot_market_usdc`: Raw bytes from Drift's USDC spot market account containing
///   current market state and interest rates.
/// - `data_user_account`: Raw bytes from Drift user account that holds the protocol's USDC
///   deposits and accumulated yield.
/// - `receipt_amount`: Amount of USDC+ tokens to redeem, scaled by 10^6 (1 USDC+ = 1_000_000)
/// 
/// ## Returns
/// Amount of USDC tokens the user will receive (with 6 decimals places)
/// 
/// ## Process
/// 1. Extract from accounts following data:
///     - AutoCompound state from controller account
///     - Drift usdc spot market state
///     - Reflect's Drift user account state
///     - USDC+ supply
/// 2. Calculate amount of USDC at Drift
/// 3. Updates the AutoCompound state with latest Drift usdc holdings
/// 4. Calculates USDC tokens based on the updated exchange rate
/// 
/// ## Formula
/// usdc = receipt_amount * deposited_vault_value / effective_supply
///
/// ## Example
/// 
/// // Redeem 100 USDC+ tokens
/// let usdc = exchange_rate_receipt_input(
///     &controller_data,
///     &spot_market_data,
///     &user_account_data,
///     42_000_000  // 42 USDC+ with 6 decimals
/// )?;
/// 
pub fn exchange_rate_receipt_input(    
    data_usdc_controller: &[u8],
    data_spot_market_usdc: &[u8],
    data_user_account: &[u8],
    data_usdc_plus_mint: &[u8],
    receipt_amount: u64,     
) -> Result<u64> {    

    // Get autocompound and recipients.
    let mut auto_compound: AutoCompound = deserialise_autocompound(data_usdc_controller)?;

    // Get supply of USDC+ tokens.
    let usdc_plus_supply = get_mint_supply(data_usdc_plus_mint)?;

    // Get usdc amount in drift.
    let usdc_amount_drift = get_usdc_amount_drift_data(data_user_account, data_spot_market_usdc)?;
    
    // Update pool value to latest value from drift.
    auto_compound.update_pool(usdc_amount_drift, &vec![10_000], usdc_plus_supply)?;
    let deposited_vault_value: u64 = auto_compound.deposited_vault_value;    
    
    compute_usdc_from_tokens(receipt_amount, deposited_vault_value, usdc_plus_supply)
}

/// Calculates USDC redemption amount for USDC+ tokens using Anchor account validation.
/// 
/// This wrapper function provides account validation before calculating the redemption value,
/// ensuring the correct Reflect protocol accounts are being used.
/// 
/// ## Parameters
/// - `usdc_controller_account`: The USDC strategy controller AccountInfo, must match `ids::usdc_controller::ID`
/// - `spot_market_usdc_account`: Drift's USDC spot market AccountInfo, must match `ids::usdc_spot_market::ID`
/// - `reflect_user_account`: Reflect's Drift user AccountInfo, must match `ids::reflect_user_account_strategy_0::ID`
/// - `receipt_amount`: Amount of USDC+ tokens to redeem, scaled by 10^6 (1 USDC+ = 1_000_000)
/// 
/// ## Returns
/// Amount of USDC the user will receive (with 6 decimal places)
/// 
/// ## Security
/// - Validates all account addresses match expected protocol accounts
/// - Prevents account substitution attacks
/// - Returns `InvalidAccount` error if any account doesn't match
/// 
/// ## Example
/// ```ignore
/// let usdc = exchange_rate_receipt_accounts(
///     &ctx.accounts.usdc_controller,
///     &ctx.accounts.spot_market,
///     &ctx.accounts.drift_user,
///     42_000_000  // 42 USDC+
/// )?;
/// ```
pub fn exchange_rate_receipt_accounts(    
    usdc_controller_account: &AccountInfo,
    spot_market_usdc_account: &AccountInfo,
    reflect_user_account: &AccountInfo,
    usdc_plus_mint_account: &AccountInfo,
    receipt_amount: u64,     
) -> Result<u64> {  
    require!(usdc_controller_account.key == &ids::usdc_controller::ID, ReflectErrorCodes::InvalidUsdcControllerAccount);
    require!(spot_market_usdc_account.key == &ids::usdc_spot_market::ID, ReflectErrorCodes::InvalidUsdcSpotMarketAccount);
    require!(reflect_user_account.key == &ids::reflect_user_account_strategy_0::ID, ReflectErrorCodes::InvalidReflectUserAccount);
    exchange_rate_receipt_input(    
        &usdc_controller_account.data.borrow(),
        &spot_market_usdc_account.data.borrow(),
        &reflect_user_account.data.borrow(),
        &usdc_plus_mint_account.data.borrow(),
        receipt_amount,     
    )
}

// cargo test --lib exchange::tests -- --nocapture

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::path::PathBuf;

    const LOCAL_CONTROLLER: &str = "usdc_controller_account.bin";
    const LOCAL_SPOT_MARKET: &str = "spot_market_usdc.bin";
    const LOCAL_USER_ACCOUNT: &str = "user_account.bin";    
    const LOCAL_USDC_PLUS_MINT: &str = "usdc_plus_mint.bin";

    fn get_test_assets_dir() -> PathBuf {
        let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        path.push("./test_assets/local/");
        path
    }

    #[test]
    fn test_compute() {
        let result = compute_usdc_from_tokens(1000, 1_100_000, 1_000_000).unwrap();
        assert_eq!(result, 1100);
        let result = compute_usdc_from_tokens(1000, 1_000_000, 0);
        assert!(result.is_err());
        let result = compute_usdc_from_tokens(u64::MAX, u64::MAX, 1);
        assert!(result.is_err());
        let result = compute_tokens_from_usdc(u64::MAX, 1, u64::MAX);
        assert!(result.is_err());
    }

    /// Tests exchange rate calculations with real mainnet/local account data.
    #[test]
    fn test_exchange_rate_functions() {
        let assets_dir = get_test_assets_dir();
        let controller_data = fs::read(assets_dir.join(LOCAL_CONTROLLER)).expect(&format!("Missing: {}", LOCAL_CONTROLLER));
        let spot_data = fs::read(assets_dir.join(LOCAL_SPOT_MARKET)).expect(&format!("Missing: {}", LOCAL_SPOT_MARKET));
        let user_data: Vec<u8> = fs::read(assets_dir.join(LOCAL_USER_ACCOUNT)).expect(&format!("Missing: {}", LOCAL_USER_ACCOUNT));
        let usdc_plus_mint_data = fs::read(assets_dir.join(LOCAL_USDC_PLUS_MINT)).expect(&format!("Missing: {}", LOCAL_USDC_PLUS_MINT));

        println!("\n=== Deposit Tests ===");
        for amount in [1_000_000, 100_000_000, 1_000_000_000] {
            let tokens = exchange_rate_usdc_input(
                &controller_data,
                &spot_data,
                &user_data,
                &usdc_plus_mint_data,
                amount,
            ).expect("Failed to calculate deposit");
            let rate = tokens as f64 / amount as f64;
            println!(
                "Deposit {:.2} USDC → {:.6} USDC+ (1 USDC = {:.6} USDC+)",
                amount as f64 / 1_000_000.0,
                tokens as f64 / 1_000_000.0,
                rate
            );
        }

        println!("\n=== Redemption Tests ===");
        for amount in [1_000_000, 100_000_000, 1_000_000_000] {
            let usdc = exchange_rate_receipt_input(
                &controller_data,
                &spot_data,
                &user_data,
                &usdc_plus_mint_data,
                amount,
            ).expect("Failed to calculate redemption");
            let rate = usdc as f64 / amount as f64;
            println!(
                "Redeem {:.2} USDC+ → {:.6} USDC (1 USDC+ = {:.6} USDC)",
                amount as f64 / 1_000_000.0,
                usdc as f64 / 1_000_000.0,
                rate
            );
        }
    }

    /// Verifies USDC → USDC+ → USDC round trip maintains value (±1 unit for rounding).
    #[test]
    fn test_round_trip_conversion() {
        let assets_dir = get_test_assets_dir();

        let controller_data = fs::read(assets_dir.join(LOCAL_CONTROLLER)).expect(&format!("Missing: {}", LOCAL_CONTROLLER));
        let spot_data = fs::read(assets_dir.join(LOCAL_SPOT_MARKET)).expect(&format!("Missing: {}", LOCAL_SPOT_MARKET));
        let user_data = fs::read(assets_dir.join(LOCAL_USER_ACCOUNT)).expect(&format!("Missing: {}", LOCAL_USER_ACCOUNT));
        let mint_data = fs::read(assets_dir.join(LOCAL_USDC_PLUS_MINT)).expect(&format!("Missing: {}", LOCAL_USDC_PLUS_MINT));

        let initial_usdc = 1_000_000_000; // 1000 USDC

        let tokens = exchange_rate_usdc_input(
            &controller_data, &spot_data, &user_data, &mint_data, initial_usdc
        ).expect("Failed to calculate deposit");
        let final_usdc = exchange_rate_receipt_input(
            &controller_data, &spot_data, &user_data, &mint_data, tokens
        ).expect("Failed to calculate redemption");

        println!(
            "\nRound trip: {} USDC → {} USDC+ → {} USDC",
            initial_usdc / 1_000_000,
            tokens / 1_000_000,
            final_usdc / 1_000_000
        );
        let difference = (final_usdc as i64 - initial_usdc as i64).abs();
        assert!(difference <= 1, "Round trip loss exceeds tolerance: {} units", difference);
    }

    /// Compares rates with and without pending yield. Uses mint supply only.
    #[test]
    fn test_yield_impact_on_exchange_rate() {
        let assets_dir = get_test_assets_dir();

        let controller_data = fs::read(assets_dir.join(LOCAL_CONTROLLER)).expect(&format!("Missing: {}", LOCAL_CONTROLLER));
        let spot_data = fs::read(assets_dir.join(LOCAL_SPOT_MARKET)).expect(&format!("Missing: {}", LOCAL_SPOT_MARKET));
        let user_data = fs::read(assets_dir.join(LOCAL_USER_ACCOUNT)).expect(&format!("Missing: {}", LOCAL_USER_ACCOUNT));
        let mint_data = vec![0u8; 82];

        // Stored (on-chain) exchange rate: V / S, with S from the mint.
        let auto_compound = deserialise_autocompound(&controller_data).expect("Failed to deserialize autocompound");
        let usdc_plus_supply = get_mint_supply(&mint_data).unwrap_or(0);
        let stored_rate = if usdc_plus_supply > 0 {
            auto_compound.deposited_vault_value as f64 / usdc_plus_supply as f64
        } else {
            1.0
        };

        // Updated rate with pending yield captured.
        let test_amount = 1_000_000_000;
        let tokens = exchange_rate_usdc_input(
            &controller_data, &spot_data, &user_data, &mint_data, test_amount
        ).expect("Failed to calculate with yield");
        let updated_rate = test_amount as f64 / tokens as f64;

        println!("\nStored rate (no pending yield): 1 USDC+ = {:.6} USDC", stored_rate);
        println!("Updated rate (with pending yield): 1 USDC+ = {:.6} USDC", updated_rate);
        println!("Yield impact: {:.4}%", ((updated_rate / stored_rate) - 1.0) * 100.0);
    }

    /// Functions should error on invalid controller data.
    #[test]
    fn test_error_handling() {
        let bad_controller = vec![0u8; 100]; // Too short to contain AutoCompound
        let valid_spot = vec![0u8; 10000];
        let valid_user = vec![0u8; 10000];
        let mint_data = vec![0u8; 82];

        assert!(exchange_rate_usdc_input(&bad_controller, &valid_spot, &valid_user, &mint_data, 1_000_000).is_err());
        assert!(exchange_rate_receipt_input(&bad_controller, &valid_spot, &valid_user, &mint_data, 1_000_000).is_err());
    }

    /// After capture, PPS is fixed until new yield accrues. We use mint supply (S) everywhere.
    #[test]
    fn test_exchange_rate_progression() {
        let assets_dir = get_test_assets_dir();

        let controller_data = fs::read(assets_dir.join(LOCAL_CONTROLLER)).expect(&format!("Missing: {}", LOCAL_CONTROLLER));
        let spot_data = fs::read(assets_dir.join(LOCAL_SPOT_MARKET)).expect(&format!("Missing: {}", LOCAL_SPOT_MARKET));
        let user_data = fs::read(assets_dir.join(LOCAL_USER_ACCOUNT)).expect(&format!("Missing: {}", LOCAL_USER_ACCOUNT));
        let usdc_plus_mint_data = fs::read(assets_dir.join(LOCAL_USDC_PLUS_MINT)).expect(&format!("Missing: {}", LOCAL_USDC_PLUS_MINT));

        
        let mut auto_compound = deserialise_autocompound(&controller_data).expect("Failed to deserialize");
        let initial_drift_value = get_usdc_amount_drift_data(&user_data, &spot_data).expect("Failed to get drift value");
        let usdc_plus_supply = get_mint_supply(&usdc_plus_mint_data).unwrap_or(0);

        println!("\n=== Exchange Rate Progression Test ===");
        println!("Initial drift value: {:.2}", initial_drift_value as f64 / 1_000_000.0);

        // Test amounts.
        let test_usdc = 1_000_000_000;   // 1000 USDC
        let test_receipt = 1_000_000_000; // 1000 USDC+

        // Simulate multiple rounds of yield accumulation and capture.
        let mut last_deposit_rate = 0.0;
        let mut last_redeem_rate = 0.0;

        for round in 0..5 {
            println!("\n--- Round {} ---", round);

            // Simulate yield accumulation (2% per round for testing)
            let simulated_yield = (initial_drift_value as f64 * 0.02 * (round + 1) as f64) as u64;
            let current_drift_value = initial_drift_value + simulated_yield;

            println!(
                "Simulated drift value: {:.2} (yield: {:.2})",
                current_drift_value as f64 / 1_000_000.0,
                simulated_yield as f64 / 1_000_000.0
            );

            // Use SUPPLY FROM MINT (S). If zero (placeholder), the recipient path won’t mint anyway
            // with cuts = [10_000] (100% pool keep), so this is safe.
            let token_supply = usdc_plus_supply;

            // New update_pool signature
            auto_compound.update_pool(current_drift_value, &[10_000], token_supply).unwrap();

            let deposit_tokens = compute_tokens_from_usdc(
                test_usdc,
                auto_compound.deposited_vault_value,
                token_supply,
            ).unwrap();

            let redeem_usdc = compute_usdc_from_tokens(
                test_receipt,
                auto_compound.deposited_vault_value,
                token_supply,
            ).unwrap();

            let deposit_rate = deposit_tokens as f64 / test_usdc as f64;
            let redeem_rate = redeem_usdc as f64 / test_receipt as f64;

            println!("After capture:");
            println!("  1 USDC = {:.6} USDC+ (deposit rate)", deposit_rate);
            println!("  1 USDC+ = {:.6} USDC (redeem rate)", redeem_rate);
            println!("  deposited_vault_value (V): {}", auto_compound.deposited_vault_value);
            println!("  token_supply (S from mint): {}", token_supply);

            if round > 0 {
                // Deposit rate should not improve after capture.
                assert!(
                    deposit_rate <= last_deposit_rate + 0.000001,
                    "Deposit rate improved after capture: {} > {} at round {}",
                    deposit_rate, last_deposit_rate, round
                );
                // Redeem rate should not worsen after capture.
                assert!(
                    redeem_rate >= last_redeem_rate - 0.000001,
                    "Redeem rate worsened after capture: {} < {} at round {}",
                    redeem_rate, last_redeem_rate, round
                );
            }

            last_deposit_rate = deposit_rate;
            last_redeem_rate = redeem_rate;

            // Consistency across amounts
            for amount in [1_000_000, 100_000_000, 10_000_000_000] {
                let tokens = compute_tokens_from_usdc(
                    amount,
                    auto_compound.deposited_vault_value,
                    token_supply,
                ).unwrap();
                let rate = tokens as f64 / amount as f64;
                assert!(
                    (rate - deposit_rate).abs() < 0.000001,
                    "Rate inconsistent for different amounts"
                );
            }
        }

        println!("\n=== Summary ===");
        println!("Deposit rate progression (USDC+ per USDC should decrease):");
        println!("  Started at: ~1.0");
        println!("  Ended at: {:.6}", last_deposit_rate);
        println!("Redeem rate progression (USDC per USDC+ should increase):");
        println!("  Started at: ~1.0");
        println!("  Ended at: {:.6}", last_redeem_rate);
    }
}