riftenlabs-indexer/src/rpc/tvl.rs

// Copyright (C) 2025-2026 Whiterun LLC
//
// This software is licensed under the GNU Affero General Public License (AGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.html

use std::collections::HashMap;

use anyhow::Result;
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use rocket::{get, State};
use serde_json::{json, Value};
use sqlx::SqlitePool;

use crate::db::{
    cauldron::poolvisitor::{
        db_visit_pool_entries, OptionalFields, OptionalPoolFields, PoolFilters, PoolVisitor,
    },
    DB,
};
use crate::rpc::err::{db_error, CachedApiResult};
use crate::rpc::response::{cached_ok, CACHE_AGGREGATE, CACHE_IMMUTABLE};

#[derive(Default)]
pub struct TvlByTokenVisitor {
    tvl: HashMap<String, (u64, u64)>,
}

impl TvlByTokenVisitor {
    pub fn into_map(self) -> HashMap<String, (u64, u64)> {
        self.tvl
    }
}

impl PoolVisitor for TvlByTokenVisitor {
    fn optional_fields_wanted(&self) -> u64 {
        OptionalPoolFields::TokenId as u64
    }

    fn visit(&mut self, sats: u64, tokens: u64, optional_fields: OptionalFields) -> Result<bool> {
        let token_id = optional_fields.token_id.unwrap().to_lowercase();

        let entry = self.tvl.entry(token_id).or_insert((0u64, 0u64));
        entry.0 += sats;
        entry.1 += tokens;

        Ok(true)
    }
}

#[derive(Default)]
struct TvlVisitor {
    sats: u64,
    tokens: u64,
}

impl PoolVisitor for TvlVisitor {
    fn optional_fields_wanted(&self) -> u64 {
        0
    }

    fn visit(&mut self, sats: u64, tokens: u64, _optional_fields: OptionalFields) -> Result<bool> {
        self.sats += sats;
        self.tokens += tokens;

        Ok(true)
    }
}

#[derive(Default)]
struct TvlSatsOnlyVisitor {
    sats: u64,
}

impl PoolVisitor for TvlSatsOnlyVisitor {
    fn optional_fields_wanted(&self) -> u64 {
        0
    }

    fn visit(&mut self, sats: u64, _tokens: u64, _optional_fields: OptionalFields) -> Result<bool> {
        self.sats += sats;

        Ok(true)
    }
}

/// Fetch TVL for all tokens by token id (deprecated)
pub async fn deprecated_get_all_token_tvl(
    pool: &SqlitePool,
    max_timestamp: usize,
) -> Result<HashMap<String, (u64, u64)>> {
    let mut visitor = TvlByTokenVisitor::default();
    db_visit_pool_entries(
        pool,
        &mut visitor,
        PoolFilters {
            timestamp_lt: Some(max_timestamp as u64),
            timestamp_lte: None,
            timestamp_gt: None,
            timestamp_gte: None,
            token_id: None,
            owner: None,
        },
    )
    .await?;

    Ok(visitor.tvl)
}

/// Fetch sats side of TVL for all tokens
pub async fn get_total_sats_tvl(pool: &SqlitePool, max_timestamp: Option<usize>) -> Result<u64> {
    let mut visitor = TvlSatsOnlyVisitor::default();
    db_visit_pool_entries(
        pool,
        &mut visitor,
        PoolFilters {
            timestamp_lt: max_timestamp.map(|t| t as u64),
            timestamp_lte: None,
            timestamp_gt: None,
            timestamp_gte: None,
            token_id: None,
            owner: None,
        },
    )
    .await?;

    Ok(visitor.sats)
}

/// Get TVL for a single token
pub async fn get_token_tvl(
    pool: &SqlitePool,
    max_timestamp: Option<usize>,
    token_id: &str,
) -> Result<(u64, u64)> {
    let mut visitor = TvlVisitor::default();
    db_visit_pool_entries(
        pool,
        &mut visitor,
        PoolFilters {
            timestamp_lt: max_timestamp.map(|t| t as u64),
            timestamp_lte: None,
            timestamp_gt: None,
            timestamp_gte: None,
            token_id: Some(token_id.to_string()),
            owner: None,
        },
    )
    .await?;

    Ok((visitor.sats, visitor.tokens))
}

/// Status: Deprecated
/// use valuelocked with optional parameters
/// used by defilama; fix adapter before removing
#[get("/tvl/<time>")]
pub async fn deprecated_tvl(time: usize, conn: &State<DB>) -> CachedApiResult<Vec<Value>> {
    let tvl: HashMap<String, (u64, u64)> = deprecated_get_all_token_tvl(&conn.cauldron_r, time)
        .await
        .map_err(db_error)?;

    let result: Vec<Value> = tvl
        .into_par_iter()
        .map(|(token, (sats, token_amount))| {
            json!({
                "token_id": token,
                "satoshis": sats,
                "token_amount": token_amount,
            })
        })
        .collect();

    Ok(cached_ok(result, CACHE_AGGREGATE))
}

/// Gives total satoshis locked for all tokens.
/// Status: Stable
///
/// - time: Unix timestamp (optional)
///
/// **Response Example:**
///
/// ```json
/// {
///     "satoshis": 1459676788
/// }
/// ```
#[get("/valuelocked?<time>")]
pub async fn valuelocked_all(time: Option<usize>, conn: &State<DB>) -> CachedApiResult<Value> {
    let sats: u64 = get_total_sats_tvl(&conn.cauldron_r, time)
        .await
        .map_err(db_error)?;

    let cache_duration = if time.is_some() {
        CACHE_IMMUTABLE
    } else {
        CACHE_AGGREGATE
    };

    Ok(cached_ok(
        json!({
            "satoshis": sats
        }),
        cache_duration,
    ))
}

/// Gives total value locked for a single token.
/// Status: Stable
///
/// - token: Token identifier / category.
/// - time: Unix timestamp (optional)
///
/// **Response Example:**
///
/// ```json
/// {
///     "satoshis": 1459676788,
///     "token_amount": 19,
///     "token_id": "b79bfc8246b5fc4707e7c7dedcb6619ef1ab91f494a790c20b0f4c422ed95b92"
/// }
/// ```
#[get("/valuelocked/<token>?<time>")]
pub async fn valuelocked_token(
    token: &str,
    time: Option<usize>,
    conn: &State<DB>,
) -> CachedApiResult<Value> {
    let (sats, token_amount) = get_token_tvl(&conn.cauldron_r, time, token)
        .await
        .map_err(db_error)?;

    let cache_duration = if time.is_some() {
        CACHE_IMMUTABLE
    } else {
        CACHE_AGGREGATE
    };

    Ok(cached_ok(
        json!({
            "token_amount": token_amount,
            "satoshis": sats
        }),
        cache_duration,
    ))
}

#[cfg(test)]
pub mod tests {
    use super::*;
    use crate::db::cauldron::{pool, tx, utxo_funding, utxo_spending};

    use crate::utiltest::mock_db_pool;
    use bitcoin_hashes::Hash;
    use bitcoincash::{BlockHash, PubkeyHash, TokenID, Txid};
    use riftenlabs_defi::cauldron::ParsedContract;
    use riftenlabs_defi::chainutil::OutPointHash;
    use rocket::http::Status;
    use rocket::local::asynchronous::Client;
    use rocket::routes;

    fn dummy_cauldron(
        txid: &Txid,
        utxo: &OutPointHash,
        token: &TokenID,
        sats: u64,
        tokens: i64,
        pkh: &PubkeyHash,
    ) -> ParsedContract {
        ParsedContract {
            pkh: *pkh,
            is_withdrawn: false,
            spent_utxo_hash: OutPointHash::all_zeros(),
            new_utxo_hash: Some(*utxo),
            new_utxo_txid: Some(*txid),
            new_utxo_n: Some(0),
            token_id: Some(*token),
            sats: Some(sats),
            token_amount: Some(tokens),
        }
    }

    async fn setup_mock_db(pool: sqlx::SqlitePool) {
        // Create required tables
        utxo_funding::create_table(&pool).await;
        utxo_spending::create_table(&pool).await;
        tx::create_table(&pool).await;
        pool::create_table(&pool).await;
        pool::dummy_init_seq();

        let mut conn = pool.acquire().await.unwrap();

        let token_zero = TokenID::all_zeros();
        let pkh_zero = PubkeyHash::all_zeros();
        let block_zero = BlockHash::all_zeros();

        // Create test data with different timestamps
        let txid1 = Txid::from_inner([0xf1; 32]);
        let txid2 = Txid::from_inner([0xf2; 32]);
        let txid3 = Txid::from_inner([0xf3; 32]);
        let utxo1 = OutPointHash::from_inner([0xe1; 32]);
        let utxo2 = OutPointHash::from_inner([0xe2; 32]);
        let utxo3 = OutPointHash::from_inner([0xe3; 32]);

        // Insert transactions with different timestamps
        tx::insert_block_tx(&mut *conn, &txid1, &block_zero, 1000)
            .await
            .unwrap();
        tx::insert_block_tx(&mut *conn, &txid2, &block_zero, 2000)
            .await
            .unwrap();
        tx::insert_block_tx(&mut *conn, &txid3, &block_zero, 3000)
            .await
            .unwrap();

        // Create cauldrons with different sats amounts
        let cauldron1 = dummy_cauldron(&txid1, &utxo1, &token_zero, 100_000, 1_000, &pkh_zero);
        let cauldron2 = dummy_cauldron(&txid2, &utxo2, &token_zero, 200_000, 2_000, &pkh_zero);
        let cauldron3 = dummy_cauldron(&txid3, &utxo3, &token_zero, 300_000, 3_000, &pkh_zero);

        // Insert utxo funding data
        utxo_funding::insert_utxo_funding(&mut *conn, &vec![cauldron1.clone()], &txid1, true)
            .await
            .unwrap();
        utxo_funding::insert_utxo_funding(&mut *conn, &vec![cauldron2.clone()], &txid2, true)
            .await
            .unwrap();
        utxo_funding::insert_utxo_funding(&mut *conn, &vec![cauldron3.clone()], &txid3, true)
            .await
            .unwrap();

        // Insert pool data
        pool::insert_new_pool(&mut *conn, &cauldron1).await.unwrap();
        pool::insert_new_pool(&mut *conn, &cauldron2).await.unwrap();
        pool::insert_new_pool(&mut *conn, &cauldron3).await.unwrap();

        // Insert pool history entries
        pool::insert_pool_history_entry(
            &mut *conn,
            &utxo1,
            &cauldron1,
            Some(1000),
            Some(1000),
            0,
            0,
        )
        .await
        .unwrap();
        pool::insert_pool_history_entry(
            &mut *conn,
            &utxo2,
            &cauldron2,
            Some(2000),
            Some(2000),
            0,
            0,
        )
        .await
        .unwrap();
        pool::insert_pool_history_entry(
            &mut *conn,
            &utxo3,
            &cauldron3,
            Some(3000),
            Some(3000),
            0,
            0,
        )
        .await
        .unwrap();
    }

    #[rocket::async_test]
    async fn test_get_valuelocked_all() {
        // Set up a fresh mock DB
        let mock_db = mock_db_pool(setup_mock_db).await;

        // Test 1: Get total TVL without time filter (should include all data)
        let result = get_total_sats_tvl(&mock_db.cauldron_r, None).await.unwrap();
        assert_eq!(result, 600_000); // 100k + 200k + 300k

        // Test 2: Get TVL with time filter (should only include data before timestamp 2500)
        let result = get_total_sats_tvl(&mock_db.cauldron_r, Some(2500))
            .await
            .unwrap();
        assert_eq!(result, 300_000); // Only 100k + 200k (before timestamp 2500)

        // Test 3: Get TVL with time filter (should only include data before timestamp 1500)
        let result = get_total_sats_tvl(&mock_db.cauldron_r, Some(1500))
            .await
            .unwrap();
        assert_eq!(result, 100_000); // Only 100k (before timestamp 1500)

        // Test 4: Get TVL with time filter (should include no data before timestamp 500)
        let result = get_total_sats_tvl(&mock_db.cauldron_r, Some(500))
            .await
            .unwrap();
        assert_eq!(result, 0); // No data before timestamp 500

        // Now test the HTTP endpoint
        let rocket = rocket::build()
            .manage(mock_db)
            .mount("/api", routes![super::valuelocked_all]);
        let client = Client::tracked(rocket)
            .await
            .expect("valid rocket instance");

        // Test HTTP endpoint without time filter
        let response = client.get("/api/valuelocked").dispatch().await;
        assert_eq!(response.status(), Status::Ok);

        let body: Value = serde_json::from_str(&response.into_string().await.unwrap()).unwrap();
        assert_eq!(body["satoshis"], 600_000);

        // Test HTTP endpoint with time filter
        let response = client.get("/api/valuelocked?time=2500").dispatch().await;
        assert_eq!(response.status(), Status::Ok);

        let body: Value = serde_json::from_str(&response.into_string().await.unwrap()).unwrap();
        assert_eq!(body["satoshis"], 300_000);
    }

    #[rocket::async_test]
    async fn test_get_token_tvl_specific_token() {
        // Set up a fresh mock DB
        let mock_db = mock_db_pool(setup_mock_db).await;

        // Use token_zero as the token ID since that's what's in the mock data
        let token_id = "0000000000000000000000000000000000000000000000000000000000000000";

        // Test 1: Get TVL for the specific token without time filter (should include all data)
        let result = get_token_tvl(&mock_db.cauldron_r, None, token_id)
            .await
            .unwrap();
        assert_eq!(result, (600_000, 6_000)); // Total: 100k+200k+300k sats, 1k+2k+3k tokens

        // Test 2: Get TVL for the specific token with time filter (before timestamp 2500)
        let result = get_token_tvl(&mock_db.cauldron_r, Some(2500), token_id)
            .await
            .unwrap();
        assert_eq!(result, (300_000, 3_000)); // Only cauldron1 and cauldron2: 100k+200k sats, 1k+2k tokens

        // Test 3: Get TVL for the specific token with time filter (before timestamp 1500)
        let result = get_token_tvl(&mock_db.cauldron_r, Some(1500), token_id)
            .await
            .unwrap();
        assert_eq!(result, (100_000, 1_000)); // Only cauldron1: 100k sats, 1k tokens

        // Test 4: Get TVL for a non-existent token
        let token_id = "1111111111111111111111111111111111111111111111111111111111111111";
        let result = get_token_tvl(&mock_db.cauldron_r, None, token_id)
            .await
            .unwrap();
        assert_eq!(result, (0, 0)); // No data for non-existent token
    }
}