predict_fun_sdk/

execution.rs

//! Predict.fun execution pipeline: auth → sign → submit with dry-run safety guard.
//!
//! [`PredictExecutionClient`] handles the full lifecycle: JWT authentication,
//! market lookup (with cache), order preparation, EIP-712 signing, and submission.
//!
//! # Network optimization
//!
//! Market metadata (token IDs, fee rates, flags) is cached after first fetch.
//! This avoids redundant `GET /markets/{id}` on every order — critical for MM bots
//! placing 10+ orders/min on the same market.

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::{anyhow, Context, Result};
use serde_json::{json, Value};
use tokio::sync::RwLock;

use crate::api::{PredictApiClient, RawApiResponse};
use crate::order::{
    predict_limit_order_amounts, PredictCreateOrderRequest, PredictOrder, PredictOrderSigner,
    PredictOutcome, PredictSide, PredictStrategy, SignedPredictOrder, BNB_MAINNET_CHAIN_ID,
};

/// Cached market metadata — everything needed to sign orders without re-fetching.
#[derive(Debug, Clone)]
pub struct MarketMeta {
    pub market_id: i64,
    pub yes_token_id: String,
    pub no_token_id: String,
    pub fee_rate_bps: u32,
    pub is_neg_risk: bool,
    pub is_yield_bearing: bool,
}

impl MarketMeta {
    /// Get token ID for an outcome.
    pub fn token_id(&self, outcome: PredictOutcome) -> &str {
        match outcome {
            PredictOutcome::Yes => &self.yes_token_id,
            PredictOutcome::No => &self.no_token_id,
        }
    }
}

#[derive(Debug, Clone)]
pub struct PredictExecConfig {
    pub api_key: String,
    pub private_key: String,
    pub chain_id: u64,
    pub live_execution: bool,
    pub fill_or_kill: bool,
}

impl PredictExecConfig {
    pub fn from_env() -> Result<Self> {
        let api_key = std::env::var("PREDICT_API_KEY")
            .context("PREDICT_API_KEY is required for Predict execution")?;
        let private_key = std::env::var("PREDICT_PRIVATE_KEY")
            .or_else(|_| std::env::var("PREDICT_TEST_PRIVATE_KEY"))
            .context("PREDICT_PRIVATE_KEY (or PREDICT_TEST_PRIVATE_KEY) is required")?;

        let live_execution = std::env::var("PREDICT_LIVE_EXECUTION")
            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
            .unwrap_or(false);

        let fill_or_kill = std::env::var("PREDICT_FILL_OR_KILL")
            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
            .unwrap_or(true);

        let chain_id = std::env::var("PREDICT_CHAIN_ID")
            .ok()
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(BNB_MAINNET_CHAIN_ID);

        Ok(Self {
            api_key,
            private_key,
            chain_id,
            live_execution,
            fill_or_kill,
        })
    }
}

#[derive(Debug, Clone)]
pub struct PredictLimitOrderRequest {
    pub market_id: i64,
    pub outcome: PredictOutcome,
    pub side: PredictSide,
    pub price_per_share: f64,
    pub quantity: f64,
    pub strategy: PredictStrategy,
    pub slippage_bps: Option<u32>,
}

#[derive(Debug, Clone)]
pub struct PredictPreparedOrder {
    pub signed_order: SignedPredictOrder,
    pub request: PredictCreateOrderRequest,
    pub is_neg_risk: bool,
    pub is_yield_bearing: bool,
}

#[derive(Debug, Clone)]
pub struct PredictSubmitResult {
    pub prepared: PredictPreparedOrder,
    pub submitted: bool,
    pub response: Option<Value>,
    pub raw: Option<RawApiResponse>,
}

/// Execution client with market metadata cache and JWT management.
///
/// The market cache avoids redundant `GET /markets/{id}` calls — token IDs,
/// fee rates, and flags don't change during a market's lifetime.
#[derive(Clone)]
pub struct PredictExecutionClient {
    pub api: PredictApiClient,
    pub signer: PredictOrderSigner,
    pub config: PredictExecConfig,
    /// Market metadata cache — public for testing. Use `market_meta()` for normal access.
    pub market_cache: Arc<RwLock<HashMap<i64, MarketMeta>>>,
}

impl PredictExecutionClient {
    pub async fn new(config: PredictExecConfig) -> Result<Self> {
        let signer = PredictOrderSigner::from_private_key(&config.private_key, config.chain_id)?;
        let api = PredictApiClient::new_mainnet(&config.api_key)?;
        let jwt = Self::authenticate_jwt(&api, &signer).await?;
        let api = api.with_jwt(jwt);

        Ok(Self {
            api,
            signer,
            config,
            market_cache: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    pub async fn from_env() -> Result<Self> {
        let cfg = PredictExecConfig::from_env()?;
        Self::new(cfg).await
    }

    /// Re-authenticate and refresh the JWT token.
    pub async fn refresh_jwt(&mut self) -> Result<()> {
        let jwt = Self::authenticate_jwt(&self.api, &self.signer).await?;
        self.api.set_jwt(jwt);
        Ok(())
    }

    pub async fn authenticate_jwt(
        api: &PredictApiClient,
        signer: &PredictOrderSigner,
    ) -> Result<String> {
        let auth_message = api.auth_message().await.context("GET /auth/message failed")?;
        let message = auth_message
            .get("data")
            .and_then(|d| d.get("message"))
            .and_then(|m| m.as_str())
            .ok_or_else(|| anyhow!("missing data.message in auth response"))?;

        let signature = signer.sign_auth_message(message)?;
        let auth = api
            .auth(&signer.address().to_string(), message, &signature)
            .await
            .context("POST /auth failed")?;

        auth.get("data")
            .and_then(|d| d.get("token"))
            .and_then(|t| t.as_str())
            .map(str::to_string)
            .ok_or_else(|| anyhow!("missing data.token in auth response"))
    }

    // === Market metadata cache ===

    /// Get cached market metadata, fetching from API if not cached.
    pub async fn market_meta(&self, market_id: i64) -> Result<MarketMeta> {
        // Fast path: check read lock
        {
            let cache = self.market_cache.read().await;
            if let Some(meta) = cache.get(&market_id) {
                return Ok(meta.clone());
            }
        }

        // Slow path: fetch and cache
        let meta = self.fetch_market_meta(market_id).await?;
        {
            let mut cache = self.market_cache.write().await;
            cache.insert(market_id, meta.clone());
        }
        Ok(meta)
    }

    /// Force-refresh market metadata from API.
    pub async fn refresh_market_meta(&self, market_id: i64) -> Result<MarketMeta> {
        let meta = self.fetch_market_meta(market_id).await?;
        let mut cache = self.market_cache.write().await;
        cache.insert(market_id, meta.clone());
        Ok(meta)
    }

    /// Pre-warm the cache for multiple markets in parallel.
    pub async fn preload_markets(&self, market_ids: &[i64]) -> Result<()> {
        let mut tasks = Vec::new();
        for &id in market_ids {
            let client = self.clone();
            tasks.push(tokio::spawn(async move { client.market_meta(id).await }));
        }
        for task in tasks {
            task.await.map_err(|e| anyhow!("join error: {}", e))??;
        }
        Ok(())
    }

    /// Clear the market metadata cache.
    pub async fn clear_cache(&self) {
        self.market_cache.write().await.clear();
    }

    fn fetch_market_meta(
        &self,
        market_id: i64,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<MarketMeta>> + Send + '_>> {
        Box::pin(async move {
            let market = self
                .api
                .get_market(market_id)
                .await
                .with_context(|| format!("GET /markets/{} failed", market_id))?;

            let data = market
                .get("data")
                .ok_or_else(|| anyhow!("missing data in market response"))?;

            let outcomes = data
                .get("outcomes")
                .and_then(|o| o.as_array())
                .ok_or_else(|| anyhow!("missing outcomes in market {}", market_id))?;

            let yes_token = extract_token_id_from_outcomes(outcomes, 1)
                .ok_or_else(|| anyhow!("missing YES token (indexSet=1) in market {}", market_id))?;
            let no_token = extract_token_id_from_outcomes(outcomes, 2)
                .ok_or_else(|| anyhow!("missing NO token (indexSet=2) in market {}", market_id))?;

            Ok(MarketMeta {
                market_id,
                yes_token_id: yes_token,
                no_token_id: no_token,
                fee_rate_bps: data
                    .get("feeRateBps")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0) as u32,
                is_neg_risk: data
                    .get("isNegRisk")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false),
                is_yield_bearing: data
                    .get("isYieldBearing")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(true),
            })
        })
    }

    // === Order operations ===

    /// Prepare a signed order using cached market metadata.
    /// First call for a market fetches from API; subsequent calls use cache.
    pub async fn prepare_limit_order(
        &self,
        req: &PredictLimitOrderRequest,
    ) -> Result<PredictPreparedOrder> {
        let meta = self.market_meta(req.market_id).await?;
        self.prepare_limit_order_with_meta(req, &meta)
    }

    /// Prepare a signed order with pre-fetched market metadata (zero network calls).
    pub fn prepare_limit_order_with_meta(
        &self,
        req: &PredictLimitOrderRequest,
        meta: &MarketMeta,
    ) -> Result<PredictPreparedOrder> {
        let token_id = meta.token_id(req.outcome);
        let price_wei = wei_from_decimal(req.price_per_share)?;
        let quantity_wei = wei_from_decimal(req.quantity)?;

        let (maker_amount, taker_amount) =
            predict_limit_order_amounts(req.side, price_wei, quantity_wei);

        let maker = self.signer.address();
        let order = PredictOrder::new_limit(
            maker,
            maker,
            token_id,
            req.side,
            maker_amount,
            taker_amount,
            meta.fee_rate_bps,
        );

        let signed_order = self
            .signer
            .sign_order(&order, meta.is_neg_risk, meta.is_yield_bearing)
            .context("failed to sign predict order")?;

        let create_request = signed_order.to_create_order_request(
            price_wei,
            req.strategy,
            req.slippage_bps,
            Some(self.config.fill_or_kill),
        );

        Ok(PredictPreparedOrder {
            signed_order,
            request: create_request,
            is_neg_risk: meta.is_neg_risk,
            is_yield_bearing: meta.is_yield_bearing,
        })
    }

    /// Submit a prepared order.
    ///
    /// - `live_execution=false`: dry-run, no POST.
    /// - `live_execution=true`: sends POST /orders.
    pub async fn submit_prepared_order(
        &self,
        prepared: PredictPreparedOrder,
    ) -> Result<PredictSubmitResult> {
        if !self.config.live_execution {
            return Ok(PredictSubmitResult {
                prepared,
                submitted: false,
                response: None,
                raw: None,
            });
        }

        let body = serde_json::to_value(&prepared.request)
            .context("failed to serialize create-order request")?;

        let raw = self
            .api
            .raw_post("/orders", &[], body, true)
            .await
            .context("POST /orders failed")?;

        let response = raw.json.clone();

        Ok(PredictSubmitResult {
            prepared,
            submitted: true,
            response,
            raw: Some(raw),
        })
    }

    /// Prepare + submit in one call (uses cached market metadata).
    pub async fn place_limit_order(
        &self,
        req: &PredictLimitOrderRequest,
    ) -> Result<PredictSubmitResult> {
        let prepared = self.prepare_limit_order(req).await?;
        self.submit_prepared_order(prepared).await
    }

    /// Remove orders from orderbook via POST /orders/remove.
    pub async fn remove_order_ids(&self, ids: &[String]) -> Result<RawApiResponse> {
        if !self.config.live_execution {
            return Ok(RawApiResponse {
                status: reqwest::StatusCode::OK,
                json: Some(json!({"success": true, "dryRun": true})),
            });
        }

        let body = json!({ "data": { "ids": ids } });
        self.api
            .raw_post("/orders/remove", &[], body, true)
            .await
            .context("POST /orders/remove failed")
    }
}

fn extract_token_id_from_outcomes(outcomes: &[Value], index_set: u64) -> Option<String> {
    outcomes
        .iter()
        .find(|o| o.get("indexSet").and_then(|v| v.as_u64()) == Some(index_set))
        .and_then(|o| o.get("onChainId"))
        .and_then(|v| v.as_str())
        .map(str::to_string)
}

fn wei_from_decimal(value: f64) -> Result<alloy_primitives::U256> {
    if !value.is_finite() || value <= 0.0 {
        return Err(anyhow!("invalid decimal value {}, expected > 0", value));
    }

    let scaled = (value * 1e18_f64).round();
    if scaled <= 0.0 {
        return Err(anyhow!("value too small after scaling: {}", value));
    }

    Ok(alloy_primitives::U256::from(scaled as u128))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_wei_from_decimal() {
        let v = wei_from_decimal(0.1).unwrap();
        assert_eq!(v.to_string(), "100000000000000000");

        let v = wei_from_decimal(1.0).unwrap();
        assert_eq!(v.to_string(), "1000000000000000000");

        assert!(wei_from_decimal(0.0).is_err());
        assert!(wei_from_decimal(-1.0).is_err());
    }

    #[test]
    fn test_extract_token_id() {
        let outcomes = vec![
            serde_json::json!({"indexSet": 1, "onChainId": "yes_token"}),
            serde_json::json!({"indexSet": 2, "onChainId": "no_token"}),
        ];

        assert_eq!(
            extract_token_id_from_outcomes(&outcomes, 1).unwrap(),
            "yes_token"
        );
        assert_eq!(
            extract_token_id_from_outcomes(&outcomes, 2).unwrap(),
            "no_token"
        );
        assert!(extract_token_id_from_outcomes(&outcomes, 3).is_none());
    }

    #[test]
    fn market_meta_token_lookup() {
        let meta = MarketMeta {
            market_id: 123,
            yes_token_id: "yes_abc".into(),
            no_token_id: "no_xyz".into(),
            fee_rate_bps: 200,
            is_neg_risk: false,
            is_yield_bearing: true,
        };
        assert_eq!(meta.token_id(PredictOutcome::Yes), "yes_abc");
        assert_eq!(meta.token_id(PredictOutcome::No), "no_xyz");
    }
}