vaea-flash-sdk 0.1.0

use solana_sdk::{
    instruction::{AccountMeta, Instruction},
    pubkey::Pubkey,
    signature::Keypair,
    signer::Signer,
    transaction::VersionedTransaction,
    message::{v0, VersionedMessage},
    commitment_config::CommitmentConfig,
};
use solana_client::nonblocking::rpc_client::RpcClient;
use std::str::FromStr;
use std::sync::Arc;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use reqwest::Client;
use thiserror::Error;
use crate::types::*;

// ═══════════════════════════════════════════════════════════
//  Error
// ═══════════════════════════════════════════════════════════

#[derive(Error, Debug)]
pub enum VaeaError {
    #[error("[{code}] {message}")]
    Protocol { code: VaeaErrorCode, message: String },

    #[error("HTTP request failed: {0}")]
    Network(#[from] reqwest::Error),

    #[error("Invalid pubkey: {0}")]
    InvalidPubkey(String),

    #[error("RPC error: {0}")]
    Rpc(String),

    #[error("Transaction failed: {0}")]
    Transaction(String),
}

impl VaeaError {
    pub fn protocol(code: VaeaErrorCode, msg: impl Into<String>) -> Self {
        Self::Protocol { code, message: msg.into() }
    }
}

// ═══════════════════════════════════════════════════════════
//  Client
// ═══════════════════════════════════════════════════════════

/// VAEA Flash — Universal Flash Loan SDK for Solana (Rust)
///
/// # Example
/// ```rust,no_run
/// use vaea_flash::{VaeaFlash, BorrowParams, VaeaConfig, Source};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     let payer = solana_sdk::signature::Keypair::new();
///     let flash = VaeaFlash::new("https://api.devnet.vaea.fi", &payer)?;
///
///     let capacity = flash.get_capacity().await?;
///     println!("Available tokens: {}", capacity.tokens.len());
///
///     let quote = flash.get_quote("SOL", 1000.0).await?;
///     println!("Fee: {}%", quote.fee_breakdown.total_fee_pct);
///     Ok(())
/// }
/// ```
pub struct VaeaFlash {
    api_url: String,
    source: Source,
    http: Client,
    rpc: Option<Arc<RpcClient>>,
    payer: Option<Arc<Keypair>>,
}

impl VaeaFlash {
    /// Create a new VaeaFlash client.
    pub fn new(api_url: &str, payer: &Keypair) -> Result<Self, VaeaError> {
        Ok(Self {
            api_url: api_url.to_string(),
            source: Source::Sdk,
            http: Client::new(),
            rpc: None,
            payer: Some(Arc::new(Keypair::try_from(payer.to_bytes().as_ref()).unwrap())),
        })
    }

    /// Create with full config including RPC for execute().
    pub fn with_rpc(api_url: &str, rpc_url: &str, payer: &Keypair) -> Result<Self, VaeaError> {
        Ok(Self {
            api_url: api_url.to_string(),
            source: Source::Sdk,
            http: Client::new(),
            rpc: Some(Arc::new(RpcClient::new(rpc_url.to_string()))),
            payer: Some(Arc::new(Keypair::try_from(payer.to_bytes().as_ref()).unwrap())),
        })
    }

    /// Create a read-only client (no wallet, no execute).
    pub fn read_only(api_url: &str) -> Self {
        Self {
            api_url: api_url.to_string(),
            source: Source::Sdk,
            http: Client::new(),
            rpc: None,
            payer: None,
        }
    }

    /// Set fee source (sdk or ui).
    pub fn with_source(mut self, source: Source) -> Self {
        self.source = source;
        self
    }

    // ═══════════════════════════════════════════════════════
    //  API methods
    // ═══════════════════════════════════════════════════════

    /// Get real-time capacity for all dynamically discovered tokens (120+).
    pub async fn get_capacity(&self) -> Result<CapacityResponse, VaeaError> {
        self.api_get("/v1/capacity").await
    }

    /// Get a detailed quote with fee breakdown.
    pub async fn get_quote(&self, token: &str, amount: f64) -> Result<QuoteResponse, VaeaError> {
        if amount <= 0.0 {
            return Err(VaeaError::protocol(VaeaErrorCode::InvalidAmount, "Amount must be > 0"));
        }
        let path = format!(
            "/v1/quote?token={}&amount={}&source={}",
            token, amount, self.source
        );
        self.api_get(&path).await
    }

    /// Build prefix + suffix instructions for a flash loan.
    pub async fn build(&self, request: &BuildRequest) -> Result<BuildResponse, VaeaError> {
        self.api_post("/v1/build", request).await
    }

    /// Check system health.
    pub async fn get_health(&self) -> Result<HealthResponse, VaeaError> {
        self.api_get("/v1/health").await
    }

    /// Get the full liquidity matrix — per-protocol per-token available liquidity.
    pub async fn get_matrix(&self) -> Result<MatrixResponse, VaeaError> {
        self.api_get("/v1/matrix").await
    }

    /// Get discovery summary — how many tokens were found during boot-time scanning.
    pub async fn get_discovery(&self) -> Result<DiscoverySummary, VaeaError> {
        self.api_get("/v1/discovery").await
    }

    /// Resolve the optimal flash loan route for any supported token.
    ///
    /// The Smart Router evaluates all direct routes across Marginfi, Kamino,
    /// and Jupiter Lend, then returns the cheapest feasible path.
    pub async fn get_route(
        &self,
        token: &str,
        amount: f64,
        max_fee_bps: u16,
    ) -> Result<ResolvedRoute, VaeaError> {
        if amount <= 0.0 {
            return Err(VaeaError::protocol(VaeaErrorCode::InvalidAmount, "Amount must be > 0"));
        }
        let path = format!(
            "/v1/vte?token={}&amount={}&source={}&max_fee_bps={}&alternatives=true",
            token, amount, self.source, max_fee_bps,
        );
        self.api_get(&path).await
    }

    /// Get information about all supported lending protocols.
    pub async fn get_sources(&self) -> Result<SourcesResponse, VaeaError> {
        self.api_get("/v1/sources").await
    }

    /// Get capacity breakdown per source (Marginfi, Kamino, Jupiter Lend) for each token.
    pub async fn get_aggregated_capacity(&self) -> Result<AggregatedCapacityResponse, VaeaError> {
        self.api_get("/v1/capacity/aggregated").await
    }

    /// Build flash loan instructions with user instructions sandwiched.
    pub async fn borrow(&self, params: &BorrowParams) -> Result<Vec<Instruction>, VaeaError> {
        let payer = self.payer.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "Payer keypair required"))?;

        // Fee guard — ONLY fetch quote if user set max_fee_bps (saves ~100ms for bots)
        if let Some(max_bps) = params.max_fee_bps {
            let quote = self.get_quote(&params.token, params.amount).await?;
            let actual_bps = (quote.fee_breakdown.total_fee_pct * 100.0) as u16;
            if actual_bps > max_bps {
                return Err(VaeaError::protocol(
                    VaeaErrorCode::FeeTooHigh,
                    format!("Fee {} bps exceeds max {} bps", actual_bps, max_bps),
                ));
            }
        }

        let request = BuildRequest {
            token: params.token.clone(),
            amount: params.amount,
            user_pubkey: payer.pubkey().to_string(),
            source: Some(self.source.to_string()),
            slippage_bps: params.slippage_bps,
            max_fee_bps: params.max_fee_bps,
        };

        let build = self.build(&request).await?;
        let mut all_ixs = Vec::new();

        for api_ix in &build.prefix_instructions {
            all_ixs.push(Self::parse_api_instruction(api_ix)?);
        }
        for ix in &params.instructions {
            all_ixs.push(ix.clone());
        }
        for api_ix in &build.suffix_instructions {
            all_ixs.push(Self::parse_api_instruction(api_ix)?);
        }

        Ok(all_ixs)
    }

    /// Build, sign, and send a flash loan transaction.
    /// Automatically uses the VAEA Address Lookup Table for TX compression.
    pub async fn execute(&self, params: BorrowParams) -> Result<String, VaeaError> {
        let rpc = self.rpc.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "RPC client required for execute(). Use VaeaFlash::with_rpc()"))?;
        let payer = self.payer.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "Payer keypair required"))?;

        let all_ixs = self.borrow(&params).await?;

        let blockhash = rpc.get_latest_blockhash_with_commitment(CommitmentConfig::confirmed())
            .await
            .map_err(|e| VaeaError::Rpc(e.to_string()))?
            .0;

        // Fetch our pre-loaded ALT for TX compression (~124 bytes saved)
        let lookup_tables = self.fetch_lookup_table(rpc).await;

        let msg = v0::Message::try_compile(
            &payer.pubkey(),
            &all_ixs,
            &lookup_tables,
            blockhash,
        ).map_err(|e| VaeaError::Transaction(e.to_string()))?;

        let tx = VersionedTransaction::try_new(VersionedMessage::V0(msg), &[payer.as_ref()])
            .map_err(|e| VaeaError::Transaction(e.to_string()))?;

        let sig = rpc.send_and_confirm_transaction(&tx)
            .await
            .map_err(|e| VaeaError::Transaction(e.to_string()))?;

        Ok(sig.to_string())
    }

    // ═══════════════════════════════════════════════════════
    //  simulate() — Dry Run
    // ═══════════════════════════════════════════════════════

    /// Simulate a flash loan transaction without sending it.
    /// Returns success/failure, CU consumption, and program logs.
    pub async fn simulate(&self, params: &BorrowParams) -> Result<SimulateResult, VaeaError> {
        let rpc = self.rpc.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "RPC required for simulate()"))?;
        let payer = self.payer.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "Payer required for simulate()"))?;

        let all_ixs = self.borrow(params).await?;
        let blockhash = rpc.get_latest_blockhash_with_commitment(CommitmentConfig::confirmed())
            .await
            .map_err(|e| VaeaError::Rpc(e.to_string()))?
            .0;

        let lookup_tables = self.fetch_lookup_table(rpc).await;
        let msg = v0::Message::try_compile(
            &payer.pubkey(),
            &all_ixs,
            &lookup_tables,
            blockhash,
        ).map_err(|e| VaeaError::Transaction(e.to_string()))?;

        let tx = VersionedTransaction::try_new(VersionedMessage::V0(msg), &[payer.as_ref()])
            .map_err(|e| VaeaError::Transaction(e.to_string()))?;

        let sim = rpc.simulate_transaction(&tx)
            .await
            .map_err(|e| VaeaError::Rpc(e.to_string()))?;

        Ok(SimulateResult {
            success: sim.value.err.is_none(),
            error: sim.value.err.map(|e| format!("{:?}", e)),
            compute_units: sim.value.units_consumed.unwrap_or(0),
            logs: sim.value.logs.unwrap_or_default(),
        })
    }

    // ═══════════════════════════════════════════════════════
    //  borrow_multi() — Multi-Token Atomic Flash Loans
    // ═══════════════════════════════════════════════════════

    /// Build a multi-token atomic flash loan with nested sandwich pattern:
    ///   prefix_A → prefix_B → [user IXs] → suffix_B → suffix_A
    pub async fn borrow_multi(&self, params: &BorrowMultiParams) -> Result<Vec<Instruction>, VaeaError> {
        let payer = self.payer.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "Payer required for borrow_multi()"))?;

        // Build prefix/suffix for each loan
        let mut all_builds = Vec::new();
        for loan in &params.loans {
            let request = BuildRequest {
                token: loan.token.clone(),
                amount: loan.amount,
                user_pubkey: payer.pubkey().to_string(),
                source: Some(self.source.to_string()),
                slippage_bps: params.slippage_bps,
                max_fee_bps: params.max_fee_bps,
            };
            all_builds.push(self.build(&request).await?);
        }

        let mut all_ixs = Vec::new();

        // All prefixes in order
        for build in &all_builds {
            for api_ix in &build.prefix_instructions {
                all_ixs.push(Self::parse_api_instruction(api_ix)?);
            }
        }

        // User instructions
        for ix in &params.instructions {
            all_ixs.push(ix.clone());
        }

        // All suffixes in reverse order (nested sandwich)
        for build in all_builds.iter().rev() {
            for api_ix in &build.suffix_instructions {
                all_ixs.push(Self::parse_api_instruction(api_ix)?);
            }
        }

        Ok(all_ixs)
    }

    // ═══════════════════════════════════════════════════════
    //  is_profitable() — Profitability Check
    // ═══════════════════════════════════════════════════════

    /// Check if a flash loan strategy is profitable after all fees.
    pub async fn is_profitable(
        &self,
        params: &crate::profitability::ProfitabilityParams,
    ) -> Result<crate::profitability::ProfitabilityResult, VaeaError> {
        let quote = self.get_quote(&params.token, params.amount).await?;
        Ok(crate::profitability::calculate_profitability(&quote, params))
    }

    // ═══════════════════════════════════════════════════════
    //  borrow_local() — Zero HTTP Instruction Building
    // ═══════════════════════════════════════════════════════

    /// Build flash loan instructions 100% locally — NO API call.
    ///
    /// ~0.1ms vs ~80ms for borrow(). Use this for latency-critical bots.
    /// The instructions are identical to what /v1/build returns.
    ///
    /// ```rust,no_run
    /// # use vaea_flash_sdk::*;
    /// # async fn example(flash: &VaeaFlash) -> Result<(), Box<dyn std::error::Error>> {
    /// let ixs = flash.borrow_local(&BorrowParams {
    ///     token: "SOL".to_string(),
    ///     amount: 1000.0,
    ///     instructions: vec![/* your arb IXs */],
    ///     slippage_bps: None,
    ///     max_fee_bps: None,
    /// })?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn borrow_local(&self, params: &BorrowParams) -> Result<Vec<Instruction>, VaeaError> {
        let payer = self.payer.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "Payer required for borrow_local()"))?;

        let tier = match self.source {
            Source::Sdk => crate::types::FlashTier::Sdk,
            Source::Ui => crate::types::FlashTier::Ui,
            Source::Protocol => crate::types::FlashTier::Protocol,
        };

        let result = crate::local_builder::local_build(crate::local_builder::LocalBuildParams {
            payer: payer.pubkey(),
            token: crate::local_builder::TokenId::Symbol(params.token.clone()),
            amount: params.amount,
            tier,
        }).map_err(|e| VaeaError::protocol(VaeaErrorCode::ApiError, e))?;

        let mut all_ixs = vec![result.begin_flash];
        all_ixs.extend(params.instructions.iter().cloned());
        all_ixs.push(result.end_flash);

        Ok(all_ixs)
    }

    // ═══════════════════════════════════════════════════════
    //  execute_local() — Zero HTTP Execution
    // ═══════════════════════════════════════════════════════

    /// Build, sign, and send a flash loan in ~100ms — NO API call.
    ///
    /// Uses local instruction building + direct RPC.
    /// Critical path: localBuild(<1ms) → getBlockhash(~50ms) → send(~50ms)
    pub async fn execute_local(&self, params: BorrowParams) -> Result<String, VaeaError> {
        let rpc = self.rpc.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "RPC required for execute_local()"))?;
        let payer = self.payer.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "Payer required for execute_local()"))?;

        let all_ixs = self.borrow_local(&params)?;

        let blockhash = rpc.get_latest_blockhash_with_commitment(CommitmentConfig::confirmed())
            .await
            .map_err(|e| VaeaError::Rpc(e.to_string()))?
            .0;

        let lookup_tables = self.fetch_lookup_table(rpc).await;

        let msg = v0::Message::try_compile(
            &payer.pubkey(),
            &all_ixs,
            &lookup_tables,
            blockhash,
        ).map_err(|e| VaeaError::Transaction(e.to_string()))?;

        let tx = VersionedTransaction::try_new(VersionedMessage::V0(msg), &[payer.as_ref()])
            .map_err(|e| VaeaError::Transaction(e.to_string()))?;

        let sig = rpc.send_and_confirm_transaction(&tx)
            .await
            .map_err(|e| VaeaError::Transaction(e.to_string()))?;

        Ok(sig.to_string())
    }

    // ═══════════════════════════════════════════════════════
    //  execute_smart() — Route-Aware Smart Execution
    // ═══════════════════════════════════════════════════════

    /// Build, sign, and send a flash loan with intelligent route selection.
    ///
    /// The Smart Router evaluates ALL available paths across Marginfi, Kamino,
    /// and Jupiter Lend, then selects the cheapest feasible route automatically.
    /// Falls back to execute_local() if the API is unavailable.
    pub async fn execute_smart(&self, params: BorrowParams) -> Result<String, VaeaError> {
        // Try smart route first
        match self.get_route(&params.token, params.amount, params.max_fee_bps.unwrap_or(0)).await {
            Ok(route) => {
                let has_feasible = route.candidates.iter().any(|c| c.feasible);
                if !has_feasible && !route.candidates.is_empty() {
                    let best = &route.candidates[0];
                    if !best.sufficient_liquidity {
                        return Err(VaeaError::protocol(
                            VaeaErrorCode::InsufficientLiquidity,
                            format!(
                                "Insufficient liquidity for {} {}. Best: {:.2} on {}.",
                                params.amount, route.token_symbol,
                                best.available_liquidity, best.protocol,
                            ),
                        ));
                    }
                }
                // Route resolved — execute via standard path
                self.execute(params).await
            }
            Err(VaeaError::Network(_)) | Err(VaeaError::Protocol { .. }) => {
                // API unavailable — fallback to local execution
                self.execute_local(params).await
            }
            Err(e) => Err(e),
        }
    }

    // ═══════════════════════════════════════════════════════
    //  read_flash_state() — Zero-CPI On-Chain State Reader
    // ═══════════════════════════════════════════════════════

    /// Read an active FlashState PDA from the chain.
    ///
    /// Use this to verify a flash loan is active without CPI — just pass
    /// the PDA as a read-only account.
    ///
    /// Cost: 1 RPC call (~5ms). On-chain equivalent: ~2K CU via vaea-flash-ctx crate.
    pub async fn read_flash_state(
        &self,
        payer_key: &Pubkey,
        token_mint: &Pubkey,
    ) -> Result<Option<FlashStateInfo>, VaeaError> {
        let rpc = self.rpc.as_ref()
            .ok_or_else(|| VaeaError::protocol(VaeaErrorCode::ApiError, "RPC required for read_flash_state()"))?;

        let program_id = Pubkey::from_str(VAEA_PROGRAM_ID)
            .map_err(|_| VaeaError::InvalidPubkey(VAEA_PROGRAM_ID.to_string()))?;

        let (flash_state_pda, _) = Pubkey::find_program_address(
            &[b"flash", payer_key.as_ref(), token_mint.as_ref()],
            &program_id,
        );

        let account_info = match rpc.get_account(&flash_state_pda).await {
            Ok(acc) => acc,
            Err(_) => return Ok(None),
        };

        if account_info.data.len() < 99 {
            return Ok(None);
        }
        if account_info.owner != program_id {
            return Ok(None);
        }

        let data = &account_info.data;
        // Skip 8-byte Anchor discriminator
        let payer = Pubkey::try_from(&data[8..40])
            .map_err(|_| VaeaError::protocol(VaeaErrorCode::ApiError, "Invalid payer in FlashState"))?;
        let mint = Pubkey::try_from(&data[40..72])
            .map_err(|_| VaeaError::protocol(VaeaErrorCode::ApiError, "Invalid mint in FlashState"))?;
        let amount = u64::from_le_bytes(data[72..80].try_into().unwrap());
        let fee = u64::from_le_bytes(data[80..88].try_into().unwrap());
        let source_tier = data[88];
        let slot_created = u64::from_le_bytes(data[89..97].try_into().unwrap());
        let bump = data[97];

        let tier = FlashTier::from_u8(source_tier).unwrap_or(FlashTier::Sdk);

        Ok(Some(FlashStateInfo {
            payer,
            token_mint: mint,
            amount,
            fee,
            source_tier,
            tier,
            slot_created,
            bump,
        }))
    }

    /// Fetch the VAEA Address Lookup Table for TX compression.
    async fn fetch_lookup_table(&self, rpc: &RpcClient) -> Vec<solana_sdk::address_lookup_table::AddressLookupTableAccount> {
        use solana_sdk::address_lookup_table::AddressLookupTableAccount;
        use crate::types::VAEA_LOOKUP_TABLE;

        match rpc.get_account(&VAEA_LOOKUP_TABLE).await {
            Ok(account) => {
                match solana_sdk::address_lookup_table::state::AddressLookupTable::deserialize(&account.data) {
                    Ok(table) => vec![AddressLookupTableAccount {
                        key: VAEA_LOOKUP_TABLE,
                        addresses: table.addresses.to_vec(),
                    }],
                    Err(_) => vec![],
                }
            }
            Err(_) => vec![], // Graceful fallback: TX works without ALT
        }
    }

    // ═══════════════════════════════════════════════════════
    //  Helpers
    // ═══════════════════════════════════════════════════════

    fn parse_api_instruction(api_ix: &ApiInstructionData) -> Result<Instruction, VaeaError> {
        let program_id = Pubkey::from_str(&api_ix.program_id)
            .map_err(|_| VaeaError::InvalidPubkey(api_ix.program_id.clone()))?;

        let accounts: Vec<AccountMeta> = api_ix.accounts.iter().map(|acc| {
            let pubkey = Pubkey::from_str(&acc.pubkey)
                .unwrap_or_default();
            if acc.is_writable {
                AccountMeta::new(pubkey, acc.is_signer)
            } else {
                AccountMeta::new_readonly(pubkey, acc.is_signer)
            }
        }).collect();

        let data = STANDARD.decode(&api_ix.data)
            .map_err(|e| VaeaError::protocol(VaeaErrorCode::ApiError, format!("Base64 decode: {}", e)))?;

        Ok(Instruction { program_id, accounts, data })
    }

    /// GET request with proper error handling
    async fn api_get<T: serde::de::DeserializeOwned>(&self, path: &str) -> Result<T, VaeaError> {
        let url = format!("{}{}", self.api_url, path);
        let res = self.http.get(&url).send().await?;
        if !res.status().is_success() {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            return Err(VaeaError::protocol(
                VaeaErrorCode::ApiError,
                format!("API returned HTTP {}: {}", status, body),
            ));
        }
        Ok(res.json().await?)
    }

    /// POST request with proper error handling
    async fn api_post<T: serde::de::DeserializeOwned, B: serde::Serialize>(&self, path: &str, body: &B) -> Result<T, VaeaError> {
        let url = format!("{}{}", self.api_url, path);
        let res = self.http.post(&url).json(body).send().await?;
        if !res.status().is_success() {
            let status = res.status();
            let body = res.text().await.unwrap_or_default();
            return Err(VaeaError::protocol(
                VaeaErrorCode::ApiError,
                format!("API returned HTTP {}: {}", status, body),
            ));
        }
        Ok(res.json().await?)
    }
}