pyra-instructions 0.4.4

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 {
    /// Asset IDs whose spot market accounts should be writable.
    pub writable_asset_ids: Vec<pyra_tokens::AssetId>,
    /// Asset IDs whose spot market accounts should be read-only
    /// (typically includes the quote market, USDC = 0).
    pub readable_asset_ids: Vec<pyra_tokens::AssetId>,
    /// Oracle pubkey for each asset ID involved.
    /// The caller fetches these from on-chain SpotMarketAccount data.
    pub oracle_pubkeys: Vec<(pyra_tokens::AssetId, Pubkey)>,
}

#[derive(Debug, PartialEq, Eq)]
pub enum RemainingAccountsError {
    /// An oracle pubkey was not provided for the given asset ID.
    MissingOracle(pyra_tokens::AssetId),
    /// The asset ID has no Drift market index.
    NoDriftMarket(pyra_tokens::AssetId),
}

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

/// 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(asset_id))` if an oracle is not provided for a
/// required asset ID.
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_asset_ids: Vec<pyra_tokens::AssetId> = Vec::new();
    for &idx in &params.writable_asset_ids {
        if !seen_asset_ids.contains(&idx) {
            seen_asset_ids.push(idx);
        }
    }
    for &idx in &params.readable_asset_ids {
        if !seen_asset_ids.contains(&idx) {
            seen_asset_ids.push(idx);
        }
    }

    // 1. Oracle accounts — one per unique market index, in order.
    for &idx in &seen_asset_ids {
        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_asset_ids.
    for &idx in &seen_asset_ids {
        let is_writable = params.writable_asset_ids.contains(&idx);
        accounts.push(AccountMeta {
            pubkey: pyra_accounts::get_drift_spot_market(idx)
                .ok_or(RemainingAccountsError::NoDriftMarket(idx))?,
            is_signer: false,
            is_writable,
        });
    }

    // 3. Token-2022 mints — for writable markets that use Token-2022,
    //    using seen_asset_ids to stay consistent with the dedup logic.
    for &idx in &seen_asset_ids {
        if !params.writable_asset_ids.contains(&idx) {
            continue;
        }
        let token = idx.token();
        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::allow_attributes,
    clippy::allow_attributes_without_reason,
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects,
    reason = "test code"
)]
mod tests {
    use super::*;
    use pyra_tokens::AssetId;
    use solana_pubkey::pubkey;

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

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_asset_ids: vec![AssetId::new(1).unwrap()],
            readable_asset_ids: vec![AssetId::new(0).unwrap()],
            oracle_pubkeys: vec![(AssetId::new(1).unwrap(), sol_oracle), (AssetId::new(0).unwrap(), 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(AssetId::new(1).unwrap()).unwrap();
        let usdc_spot_market = pyra_accounts::get_drift_spot_market(AssetId::new(0).unwrap()).unwrap();
        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 (asset_id 6) uses Token-2022
        let pyusd_oracle = pubkey!("BAtFj4kQttZRVep3UZS2aZRDixkGYgWsbqTBVDbnSsPF");
        let usdc_oracle = pubkey!("En8hkHLkRe9d9DraYmBTrus518BvmVH448YcvmrFM6Ce");

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_asset_ids: vec![AssetId::new(6).unwrap()],
            readable_asset_ids: vec![AssetId::new(0).unwrap()],
            oracle_pubkeys: vec![(AssetId::new(6).unwrap(), pyusd_oracle), (AssetId::new(0).unwrap(), 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 = AssetId::new(6).unwrap().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_asset_ids: vec![AssetId::new(0).unwrap()],
            readable_asset_ids: vec![AssetId::new(0).unwrap()],
            oracle_pubkeys: vec![(AssetId::new(0).unwrap(), 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 (asset_id 1) uses regular SPL Token — no mint appended
        let sol_oracle = pubkey!("BAtFj4kQttZRVep3UZS2aZRDixkGYgWsbqTBVDbnSsPF");

        let result = get_remaining_accounts(&RemainingAccountsParams {
            writable_asset_ids: vec![AssetId::new(1).unwrap()],
            readable_asset_ids: vec![],
            oracle_pubkeys: vec![(AssetId::new(1).unwrap(), 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_asset_ids: vec![],
            readable_asset_ids: 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_asset_ids: vec![AssetId::new(1).unwrap()],
            readable_asset_ids: vec![AssetId::new(0).unwrap()],
            oracle_pubkeys: vec![], // no oracles provided
        });

        assert_eq!(result, Err(RemainingAccountsError::MissingOracle(AssetId::new(1).unwrap())));
    }
}