riftenlabs-indexer/src/db/cauldron/pool.rs

// Copyright (C) 2024-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, VecDeque},
    sync::atomic::AtomicI64,
};

use crate::db::blob::{blob_to_display_hex, display_hex_to_blob, FromBlob, ToBlob};
use crate::def::PoolID;
use anyhow::{Context, Result};
use bitcoin_hashes::hex::ToHex;
use bitcoincash::TokenID;
use log::{debug, info, warn};
use malachite::Integer;
use riftenlabs_defi::{cauldron::ParsedContract, chainutil::OutPointHash};
use rusqlite::{params, Connection, Row};
use rust_decimal::prelude::Zero;
use serde::{Serialize, Serializer};

use crate::rpc::apy::PoolSnapshot;

// Custom serialization function for malachite::Integer
fn serialize_integer_as_string<S>(integer: &Integer, serializer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    serializer.serialize_str(&integer.to_string())
}

pub fn create_table(conn: &Connection) {
    conn.execute(
        "
        CREATE TABLE pool (
            creation_utxo BLOB PRIMARY KEY REFERENCES utxo_funding(new_utxo_hash) ON DELETE CASCADE,
            owner_pkh BLOB NOT NULL,
            token_id BLOB NOT NULL,
            withdrawn_in_utxo BLOB REFERENCES utxo_spending(spent_utxo_hash) ON DELETE SET NULL
        )",
        [],
    )
    .expect("failed to create table pool");

    conn.execute(
        "CREATE TABLE pool_history_entry (
            utxo BLOB PRIMARY KEY REFERENCES utxo_funding(new_utxo_hash) ON DELETE CASCADE,
            pool BLOB REFERENCES pool(creation_utxo) ON DELETE CASCADE,
            token_id BLOB NOT NULL,
            txid BLOB REFERENCES tx(txid) ON DELETE CASCADE,
            tx_pos INT NOT NULL,
            mtp_timestamp BIGINT,
            first_seen_timestamp BIGINT,
            effective_timestamp BIGINT GENERATED ALWAYS AS (COALESCE(first_seen_timestamp, mtp_timestamp)),
            sequence BIGINT NOT NULL,
            sats BIGINT NOT NULL,
            token_amount BIGINT NOT NULL,
            sats_delta BIGINT NOT NULL,
            token_delta BIGINT NOT NULL
        )",
        [],
    )
    .expect("failed to create table pool_history_entry");

    // these two indexes are for speeding up looking for pools owned by user (active rpc with pkh filter)
    conn.execute(
        "CREATE INDEX idx_pool_history_entry_pool_sequence ON pool_history_entry(pool, sequence)",
        params![],
    )
    .unwrap();
    conn.execute(
        "CREATE INDEX idx_pool_owner_pkh ON pool(owner_pkh)",
        params![],
    )
    .unwrap();

    // speed up queries that filter pools on timestamp
    conn.execute(
        "CREATE INDEX idx_pool_withdrawn_in_utxo ON pool(withdrawn_in_utxo)",
        params![],
    )
    .unwrap();

    // speeds up pool visitor query
    conn.execute(
        "CREATE INDEX idx_pool_history_entry_pool_timestamp_sequence ON pool_history_entry (
        pool,
        effective_timestamp,
        sequence DESC)",
        params![],
    )
    .unwrap();

    // for pool visitor; filtering on token
    conn.execute(
        "CREATE INDEX idx_pool_history_entry_token_id ON pool_history_entry(token_id)",
        params![],
    )
    .unwrap();
}

fn get_pool_by_utxo(conn: &Connection, utxo_hash: &OutPointHash) -> Result<Option<OutPointHash>> {
    let mut stmt = conn.prepare("SELECT pool FROM pool_history_entry WHERE utxo = ?")?;

    let mut row = stmt.query([utxo_hash.to_blob()])?;
    let utxo_blob: Option<Vec<u8>> = row.next()?.map(|r| r.get(0).unwrap());

    match utxo_blob {
        Some(blob) => Ok(Some(
            OutPointHash::from_blob(&blob).expect("invalid original_utxo utxo in db"),
        )),
        None => Ok(None),
    }
}

pub fn flag_as_withdrawn(
    conn: &Connection,
    pool_utxo: &OutPointHash,
    cauldron: &ParsedContract,
) -> Result<()> {
    conn.execute(
        "UPDATE pool SET withdrawn_in_utxo = ? WHERE creation_utxo = ?",
        params![cauldron.spent_utxo_hash.to_blob(), pool_utxo.to_blob()],
    )
    .map_err(|e| anyhow::anyhow!("failed flag pool as withdrawn. Original error: {:?}", e))?;

    Ok(())
}

pub fn insert_new_pool(conn: &Connection, cauldron: &ParsedContract) -> Result<()> {
    // or replace, as it could have been added in mempool, then block
    conn.execute(
        "INSERT OR REPLACE INTO pool (creation_utxo, owner_pkh, token_id, withdrawn_in_utxo) VALUES (?, ?, ?, ?)",
        params![
            cauldron.new_utxo_hash.expect("outpoint hash for new pool missing").to_blob(),
            cauldron.pkh.to_blob(),
            cauldron.token_id.expect("token id for new pool missing").to_blob(),
            None::<Vec<u8>>,
        ]
    ).map_err(|e| {
        anyhow::anyhow!(
            "failed to insert new pool. Original error: {:?}",
            e
        )
    })?;

    Ok(())
}

/// Next sequence number in the `pool_history_entry` table
static NEXT_SEQUENCE: AtomicI64 = AtomicI64::new(-10);

pub fn initialize_seq(conn: &Connection) {
    let mut s = conn
        .prepare("SELECT IFNULL(MAX(sequence), 0) + 1 FROM pool_history_entry")
        .unwrap();
    let mut q = s.query(params![]).unwrap();
    let seq: i64 = q.next().unwrap().unwrap().get(0).unwrap();
    NEXT_SEQUENCE.store(seq, std::sync::atomic::Ordering::SeqCst);
}

#[allow(dead_code)] // for unit tests
pub fn dummy_init_seq() {
    if NEXT_SEQUENCE.load(std::sync::atomic::Ordering::SeqCst) < 0 {
        NEXT_SEQUENCE.store(42, std::sync::atomic::Ordering::SeqCst);
    }
}

pub fn insert_pool_history_entry(
    conn: &Connection,
    pool: &OutPointHash,
    cauldron: &ParsedContract,
    mtp_timestamp: Option<u64>,
    first_seen_timestamp: Option<u64>,
    sats_delta: i64,
    token_delta: i64,
) -> Result<()> {
    let next_seq = NEXT_SEQUENCE.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
    assert!(next_seq >= 0);

    conn.execute(
        "INSERT INTO pool_history_entry (utxo, pool, token_id, txid, tx_pos, mtp_timestamp, first_seen_timestamp, sequence, sats, token_amount, sats_delta, token_delta)
         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
         ON CONFLICT(utxo) DO UPDATE SET
            pool = excluded.pool,
            token_id = excluded.token_id,
            txid = excluded.txid,
            tx_pos = excluded.tx_pos,
            mtp_timestamp = COALESCE(excluded.mtp_timestamp, pool_history_entry.mtp_timestamp),
            first_seen_timestamp = COALESCE(excluded.first_seen_timestamp, pool_history_entry.first_seen_timestamp),
            sequence = excluded.sequence",
        params![
            cauldron
                .new_utxo_hash
                .expect("utxo hash on new pool history entry")
                .to_blob(),
            pool.to_blob(),
            cauldron.token_id.expect("token id on new pool history entry").to_blob(),
            cauldron
                .new_utxo_txid
                .expect("txid of new pool history entry")
                .to_blob(),
            cauldron
                .new_utxo_n
                .expect("utxo index of new pool history entry"),
            mtp_timestamp,
            first_seen_timestamp,
            next_seq,
            cauldron.sats,
            cauldron.token_amount,
            sats_delta,
            token_delta
        ],
    )
    .map_err(|e| {
        anyhow::anyhow!(
            "failed to insert pool history entry. Original error: {:?}",
            e
        )
    })?;
    Ok(())
}

pub fn update_pool_history(
    conn: &Connection,
    cauldrons: Vec<ParsedContract>,
    mtp_timestamp: Option<u64>,
    first_seen_timestamp: Option<u64>,
) -> Result<()> {
    if cauldrons.is_empty() {
        return Ok(());
    }

    let mut queue = VecDeque::from(cauldrons);

    while let Some(current) = queue.pop_front() {
        let (pool_utxo, is_new) = match get_pool_by_utxo(conn, &current.spent_utxo_hash)? {
            Some(c) => (c, false),
            None => {
                // Not in our database. Check if child of another interaction in current batch.
                let has_parent = queue
                    .iter()
                    .any(|parent| Some(current.spent_utxo_hash) == parent.new_utxo_hash);

                if has_parent {
                    // Was child of current batch.
                    // Process later (after parent).
                    queue.push_back(current);
                    continue;
                }

                if let Some(utxo) = current.new_utxo_hash {
                    // Not in existing pool or child of a cauldron in current batch.
                    // This is a new pool.
                    assert!(!current.is_withdrawn);
                    (utxo, true)
                } else {
                    // This is a new pool that is immediately withdrawn. Just ignore.
                    assert!(current.is_withdrawn);
                    continue;
                }
            }
        };

        if current.is_withdrawn {
            info!(
                "LP {} withdrawn in {}",
                pool_utxo.to_hex(),
                current.spent_utxo_hash.to_hex()
            );
            flag_as_withdrawn(conn, &pool_utxo, &current)?;
        } else if is_new {
            info!(
                "Cauldron LP created in tx {}",
                current
                    .new_utxo_txid
                    .expect("expected txid in new LP")
                    .to_hex()
            );
            insert_new_pool(conn, &current)?;
            insert_pool_history_entry(
                conn,
                &pool_utxo,
                &current,
                mtp_timestamp,
                first_seen_timestamp,
                0, // no delta for new pool
                0, // no token delta for new pool
            )?;
        } else {
            debug!("New entry for pool {}", pool_utxo.to_hex());
            let prev_entry = get_pool_history_entry(conn, current.spent_utxo_hash)?;
            let sats_delta = if let Some(current_sats) = current.sats {
                current_sats as i64 - prev_entry.sats as i64
            } else {
                0
            };
            let token_delta = if let Some(current_token_amount) = current.token_amount {
                current_token_amount - prev_entry.token_amount as i64
            } else {
                0
            };
            insert_pool_history_entry(
                conn,
                &pool_utxo,
                &current,
                mtp_timestamp,
                first_seen_timestamp,
                sats_delta,
                token_delta,
            )?;
        }
    }

    Ok(())
}

/// Strategy for which pool state to select when two pool states are available for a timestamp
#[derive(PartialEq, Eq)]
enum SnapshotSelection {
    UseBefore,
    UseAfter,
}

/// Get pools and their state nearest to given timestamp.
fn get_nearest_entries(
    conn: &Connection,
    timestamp: i64,
    token_id: Option<&str>,
    owner_pkh: Option<&str>,
    resolution: SnapshotSelection,
) -> Result<HashMap<String, PoolSnapshot>> {
    let extra_filters = match (token_id.is_some(), owner_pkh.is_some()) {
        (true, true) => "token_id = ?2 AND owner_pkh = ?3",
        (true, false) => "token_id = ?2",
        (false, true) => "owner_pkh = ?2",
        (false, false) => "1=1",
    };

    let params = match (token_id.is_some(), owner_pkh.is_some()) {
        (true, true) => params![timestamp, token_id, owner_pkh],
        (true, false) => params![timestamp, token_id],
        (false, true) => params![timestamp, owner_pkh],
        (false, false) => params![timestamp],
    };

    let ts = "effective_timestamp";

    let preferred = match resolution {
        SnapshotSelection::UseBefore => "ts ASC", // Prefer the lower timestamp first
        SnapshotSelection::UseAfter => "ts DESC", // Prefer the higher timestamp first
    };

    let query = format!(
        "WITH NearestEntries AS (
            -- Get the closest lower or equal to the timestamp (MAX for <= timestamp)
            SELECT phe.pool, phe.sats, phe.token_amount, {ts} AS ts, phe.sequence
            FROM pool_history_entry phe
            WHERE {ts} <= ?1
            AND phe.pool IN (
                SELECT creation_utxo
                FROM pool
                WHERE {extra_filters}
            )
            GROUP BY phe.pool
            HAVING MAX({ts})

            UNION ALL

            -- Get the closest greater than the timestamp (MIN for > timestamp)
            SELECT phe.pool, phe.sats, phe.token_amount, {ts} AS ts, phe.sequence
            FROM pool_history_entry phe
            WHERE {ts} > ?1
            AND phe.pool IN (
                SELECT creation_utxo
                FROM pool
                WHERE {extra_filters}
            )
            GROUP BY phe.pool
            HAVING MIN({ts})
        )
        SELECT * FROM NearestEntries
        ORDER BY {preferred}",
    );

    let mut stmt = conn.prepare(&query)?;
    let mut rows = stmt.query(params)?;

    let from_row = |row: &Row<'_>| -> Result<PoolSnapshot> {
        Ok(PoolSnapshot {
            pool_id: row.get(0)?,
            sats: row.get(1)?,
            token_amount: row.get(2)?,
            timestamp: row.get(3)?,
        })
    };

    let mut pools: HashMap<String, PoolSnapshot> = HashMap::default();

    while let Some(row) = rows.next()? {
        let pool_snapshot = from_row(row)?;
        if !pools.contains_key(&pool_snapshot.pool_id) {
            let existed = pools.insert(pool_snapshot.pool_id.clone(), pool_snapshot.clone());
            assert!(existed.is_none());
        }
    }

    Ok(pools)
}

/// Returns
pub fn get_pool_period_snapshot(
    conn: &Connection,
    token_id: Option<&str>,
    owner_pkh: Option<&str>,
    start: i64,
    end: i64,
) -> Result<Vec<(PoolSnapshot, PoolSnapshot)>> {
    let pools_start = get_nearest_entries(
        conn,
        start,
        token_id,
        owner_pkh,
        SnapshotSelection::UseBefore,
    )?;
    let mut pools_end =
        get_nearest_entries(conn, end, token_id, owner_pkh, SnapshotSelection::UseAfter)?;

    let mut pools: Vec<(PoolSnapshot, PoolSnapshot)> = Vec::default();
    for (start_pool_id, start_pool) in pools_start {
        let end_pool = match pools_end.remove(&start_pool_id) {
            Some(end) => end,
            None => {
                warn!("Found no 'end pool' for {start_pool_id}");
                continue;
            }
        };

        let duration = end_pool.timestamp.saturating_sub(start_pool.timestamp);
        if !duration.is_zero() {
            pools.push((start_pool, end_pool))
        }
    }

    if !pools_end.is_empty() {
        warn!("Found {} end pools not in start pools", pools_end.len());
    }

    Ok(pools)
}

#[derive(Serialize)]
pub struct PoolHistoryEntry {
    txid: String,
    sats: u64,
    token_amount: u64,
    timestamp: u64,
    #[serde(serialize_with = "serialize_integer_as_string")]
    k: Integer,
}

fn get_pool_history_entry(conn: &Connection, utxo_hash: OutPointHash) -> Result<PoolHistoryEntry> {
    let mut stmt = conn.prepare("SELECT txid, sats, token_amount, effective_timestamp as timestamp FROM pool_history_entry WHERE utxo = ?")?;
    let mut rows = stmt.query(params![utxo_hash.to_blob()])?;
    let row = rows
        .next()?
        .context("no pool history entry found for UTXO hash")?;
    let txid_blob: Vec<u8> = row.get(0)?;
    let txid = blob_to_display_hex::<bitcoincash::Txid>(&txid_blob)?;
    let sats: u64 = row.get(1)?;
    let token_amount: u64 = row.get(2)?;
    Ok(PoolHistoryEntry {
        txid,
        sats,
        token_amount,
        timestamp: row.get(3)?,
        k: Integer::from(sats) * Integer::from(token_amount),
    })
}

pub fn db_pool_history(
    conn: &Connection,
    pool: &PoolID,
    start_time: u64,
) -> Result<Vec<PoolHistoryEntry>> {
    let query = "SELECT
        phe.txid,
        phe.sats,
        phe.token_amount,
        phe.effective_timestamp as timestamp
    FROM
        pool_history_entry phe
    WHERE
        phe.pool = ?1
        AND timestamp >= ?2
    ORDER BY
        phe.sequence ASC;
    ";

    let mut stmt = conn.prepare(query)?;
    let mut rows = stmt.query(params![pool.to_blob(), start_time])?;

    let mut history: Vec<PoolHistoryEntry> = Vec::default();

    while let Some(row) = rows.next()? {
        let txid_blob: Vec<u8> = row.get(0)?;
        let txid = blob_to_display_hex::<bitcoincash::Txid>(&txid_blob)?;
        let sats = row.get(1)?;
        let token_amount = row.get(2)?;

        let k = Integer::from(sats) * Integer::from(token_amount);

        history.push(PoolHistoryEntry {
            txid,
            sats,
            token_amount,
            timestamp: row.get(3)?,
            k,
        });
    }

    Ok(history)
}

pub fn db_pool_get_details(db: &Connection, pool: &PoolID) -> Result<(String, String)> {
    let (token_blob, owner_pkh_blob): (Vec<u8>, Vec<u8>) = db.query_row(
        "SELECT token_id, owner_pkh FROM pool WHERE creation_utxo = ?1",
        [pool.to_blob()],
        |row| Ok((row.get(0)?, row.get(1)?)),
    )?;
    // token_id needs byte reversal (Bitcoin hash type)
    let token_hex = blob_to_display_hex::<TokenID>(&token_blob)?;
    // owner_pkh does NOT need byte reversal (PubkeyHash)
    let owner_pkh_hex = hex::encode(&owner_pkh_blob);
    Ok((token_hex, owner_pkh_hex))
}

pub fn db_pool_id_from_utxo(conn: &Connection, utxo_hash: &OutPointHash) -> Result<Option<String>> {
    let mut stmt = conn.prepare("SELECT pool FROM pool_history_entry WHERE utxo = ?")?;
    let mut rows = stmt.query(params![utxo_hash.to_blob()])?;

    if let Some(row) = rows.next()? {
        let pool_blob: Vec<u8> = row.get(0)?;
        let pool_id = blob_to_display_hex::<PoolID>(&pool_blob)?;
        Ok(Some(pool_id))
    } else {
        Ok(None)
    }
}

/// Get total volume in satoshis across all tokens for a given time period
pub fn get_total_volume_sats(
    db: &Connection,
    start_timestamp: u64,
    end_timestamp: u64,
) -> anyhow::Result<i64> {
    let sql = "
    SELECT
        COALESCE(SUM(ABS(phe.sats_delta)), 0) AS total_volume_sats
    FROM pool_history_entry phe
    JOIN pool p ON phe.pool = p.creation_utxo
    JOIN tx ON phe.txid = tx.txid
    WHERE tx.effective_timestamp BETWEEN ? AND ?";

    let total_volume: i64 = db.query_row(sql, params![start_timestamp, end_timestamp], |row| {
        row.get(0)
    })?;

    Ok(total_volume)
}

/// Get volume in both satoshis and tokens for a specific token for a given time period
pub fn get_token_volume_sats(
    db: &Connection,
    start_timestamp: u64,
    end_timestamp: u64,
    token_id: &str,
) -> anyhow::Result<(i64, i64)> {
    // Start with tx table (smaller, indexed) then join to phe
    // This should use the new idx_tx_effective_timestamp_txid index for fast timestamp filtering
    let sql = "
    SELECT
        COALESCE(SUM(ABS(phe.sats_delta)), 0) AS token_volume_sats,
        COALESCE(SUM(ABS(phe.token_delta)), 0) AS token_volume_tokens
    FROM tx
    JOIN pool_history_entry phe ON tx.txid = phe.txid
    JOIN pool p ON phe.pool = p.creation_utxo
    WHERE tx.effective_timestamp BETWEEN ? AND ?
        AND p.token_id = ?";
    // Use proper blob conversion for Bitcoin hash types (handles byte reversal)
    let token_blob = display_hex_to_blob::<TokenID>(token_id)?;
    let (sats_volume, token_volume): (i64, i64) = db.query_row(
        sql,
        params![start_timestamp, end_timestamp, token_blob],
        |row| Ok((row.get(0)?, row.get(1)?)),
    )?;

    Ok((sats_volume, token_volume))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::cauldron::prepare_tables;
    use crate::utiltest::mock_db_pool;
    use rusqlite::Connection;

    fn setup_test_db(conn: &Connection) {
        prepare_tables(conn);
    }

    #[test]
    fn test_volume_no_trades_in_period() {
        // tests an issue where volume function wouild fail if no trades exist in the time period
        let mock_db = mock_db_pool(setup_test_db);
        let conn = mock_db.cauldron_r.get().unwrap();

        let start_timestamp = 1755508856u64;
        let end_timestamp = 1755595256u64;
        let token_id = "f6677f3d3805d70949b375d36e094ff0ec9ece2a2cb1fde6d8b0e90b368f1f63";

        let result = get_token_volume_sats(&conn, start_timestamp, end_timestamp, token_id);
        let (sats_volume, token_volume) = result.unwrap();
        assert_eq!(sats_volume, 0);
        assert_eq!(token_volume, 0);

        let result = get_total_volume_sats(&conn, start_timestamp, end_timestamp);
        let total_volume = result.unwrap();
        assert_eq!(total_volume, 0);
    }
}