pyra-instructions 0.4.0

use core::fmt;
use solana_program::instruction::AccountMeta;
use solana_program::pubkey::Pubkey;

/// Parameters for building Drift remaining accounts.
///
/// The caller is responsible for fetching oracle pubkeys from on-chain
/// `SpotMarketAccount` data and passing them in. This keeps the helper
/// synchronous and dependency-free.
pub struct RemainingAccountsParams {
    /// Market indexes whose spot market accounts should be writable.
    pub writable_spot_market_indexes: Vec<u16>,
    /// Market indexes whose spot market accounts should be read-only
    /// (typically includes the quote market, USDC = 0).
    pub readable_spot_market_indexes: Vec<u16>,
    /// Oracle pubkey for each market index involved.
    /// The caller fetches these from on-chain SpotMarketAccount data.
    pub oracle_pubkeys: Vec<(u16, Pubkey)>,
}

#[derive(Debug, PartialEq, Eq)]
pub enum RemainingAccountsError {
    /// An oracle pubkey was not provided for the given market index.
    MissingOracle(u16),
}

impl fmt::Display for RemainingAccountsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::MissingOracle(idx) => write!(f, "missing oracle for market index {idx}"),
        }
    }
}

/// Builds the `remaining_accounts` vec required by Drift instructions
/// (deposit, fulfill_withdraw, etc.).
///
/// Returns accounts in the order Drift expects:
/// 1. Oracle accounts (read-only)
/// 2. Spot market accounts (writable or read-only per the params)
/// 3. Token-2022 mint accounts (read-only, for any writable market that uses Token-2022)
///
/// Returns `Err(MissingOracle(index))` if an oracle is not provided for a
/// required market index.
pub fn get_remaining_accounts(
    params: &RemainingAccountsParams,
) -> Result<Vec<AccountMeta>, RemainingAccountsError> {
    let mut accounts = Vec::new();

    // Deduplicate and collect all unique market indexes, preserving order
    // (writable first, then readable).
    let mut seen_indexes: Vec<u16> = Vec::new();
    for &idx in &params.writable_spot_market_indexes {
        if !seen_indexes.contains(&idx) {
            seen_indexes.push(idx);
        }
    }
    for &idx in &params.readable_spot_market_indexes {
        if !seen_indexes.contains(&idx) {
            seen_indexes.push(idx);
        }
    }

    // 1. Oracle accounts — one per unique market index, in order.
    for &idx in &seen_indexes {
        let (_, oracle) = params
            .oracle_pubkeys
            .iter()
            .find(|(mi, _)| *mi == idx)
            .ok_or(RemainingAccountsError::MissingOracle(idx))?;
        accounts.push(AccountMeta {
            pubkey: *oracle,
            is_signer: false,
            is_writable: false,
        });
    }

    // 2. Spot market accounts — writable if in writable_spot_market_indexes.
    for &idx in &seen_indexes {
        let is_writable = params.writable_spot_market_indexes.contains(&idx);
        accounts.push(AccountMeta {
            pubkey: pyra_accounts::get_drift_spot_market(idx),
            is_signer: false,
            is_writable,
        });
    }

    // 3. Token-2022 mints — for writable markets that use Token-2022,
    //    using seen_indexes to stay consistent with the dedup logic.
    for &idx in &seen_indexes {
        if !params.writable_spot_market_indexes.contains(&idx) {
            continue;
        }
        if let Some(token) = pyra_tokens::Token::find_by_drift_market_index(idx) {
            if token.token_program == pyra_tokens::TOKEN_2022_PROGRAM_ID {
                accounts.push(AccountMeta {
                    pubkey: token.mint,
                    is_signer: false,
                    is_writable: false,
                });
            }
        }
    }

    Ok(accounts)
}

#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
    use super::*;
    use solana_pubkey::pubkey;

    #[test]
    fn test_single_writable_market_with_quote() {
        // Deposit wSOL (market 1), quote market USDC (market 0) readable
        let sol_oracle = pubkey!("BAtFj4kQttZRVep3UZS2aZRDixkGYgWsbqTBVDbnSsPF");
        let usdc_oracle = pubkey!("En8hkHLkRe9d9DraYmBTrus518BvmVH448YcvmrFM6Ce");

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_spot_market_indexes: vec![1],
            readable_spot_market_indexes: vec![0],
            oracle_pubkeys: vec![(1, sol_oracle), (0, usdc_oracle)],
        })
        .expect("should succeed");

        // Should have: 2 oracles + 2 spot markets = 4 accounts
        assert_eq!(result.len(), 4);

        // Oracles (writable market first, then readable)
        assert_eq!(result[0].pubkey, sol_oracle);
        assert!(!result[0].is_writable);
        assert_eq!(result[1].pubkey, usdc_oracle);
        assert!(!result[1].is_writable);

        // Spot markets
        let sol_spot_market = pyra_accounts::get_drift_spot_market(1);
        let usdc_spot_market = pyra_accounts::get_drift_spot_market(0);
        assert_eq!(result[2].pubkey, sol_spot_market);
        assert!(result[2].is_writable);
        assert_eq!(result[3].pubkey, usdc_spot_market);
        assert!(!result[3].is_writable);
    }

    #[test]
    fn test_token_2022_appends_mint() {
        // PYUSD (market 22) uses Token-2022
        let pyusd_oracle = pubkey!("BAtFj4kQttZRVep3UZS2aZRDixkGYgWsbqTBVDbnSsPF");
        let usdc_oracle = pubkey!("En8hkHLkRe9d9DraYmBTrus518BvmVH448YcvmrFM6Ce");

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_spot_market_indexes: vec![22],
            readable_spot_market_indexes: vec![0],
            oracle_pubkeys: vec![(22, pyusd_oracle), (0, usdc_oracle)],
        })
        .expect("should succeed");

        // 2 oracles + 2 spot markets + 1 Token-2022 mint = 5
        assert_eq!(result.len(), 5);

        // Last account should be PYUSD mint, read-only
        let pyusd = pyra_tokens::Token::find_by_drift_market_index(22)
            .expect("PYUSD should be a supported token");
        assert_eq!(result[4].pubkey, pyusd.mint);
        assert!(!result[4].is_writable);
        assert!(!result[4].is_signer);
    }

    #[test]
    fn test_deduplicates_overlapping_indexes() {
        // Market 0 appears in both writable and readable — should be writable, listed once
        let usdc_oracle = pubkey!("En8hkHLkRe9d9DraYmBTrus518BvmVH448YcvmrFM6Ce");

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_spot_market_indexes: vec![0],
            readable_spot_market_indexes: vec![0],
            oracle_pubkeys: vec![(0, usdc_oracle)],
        })
        .expect("should succeed");

        // 1 oracle + 1 spot market = 2
        assert_eq!(result.len(), 2);
        assert!(result[1].is_writable); // writable wins
    }

    #[test]
    fn test_spl_token_market_no_mint_appended() {
        // wSOL (market 1) uses regular SPL Token — no mint appended
        let sol_oracle = pubkey!("BAtFj4kQttZRVep3UZS2aZRDixkGYgWsbqTBVDbnSsPF");

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_spot_market_indexes: vec![1],
            readable_spot_market_indexes: vec![],
            oracle_pubkeys: vec![(1, sol_oracle)],
        })
        .expect("should succeed");

        // 1 oracle + 1 spot market = 2 (no mint)
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_empty_params() {
        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_spot_market_indexes: vec![],
            readable_spot_market_indexes: vec![],
            oracle_pubkeys: vec![],
        })
        .expect("should succeed");

        assert!(result.is_empty());
    }

    #[test]
    fn test_missing_oracle_returns_error() {
        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_spot_market_indexes: vec![1],
            readable_spot_market_indexes: vec![0],
            oracle_pubkeys: vec![], // no oracles provided
        });

        assert_eq!(result, Err(RemainingAccountsError::MissingOracle(1)));
    }
}