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

// Copyright (C) 2024 Riften Labs AS
//
// 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::def::PoolID;
use anyhow::{Context, Result};
use bitcoin_hashes::hex::{FromHex, ToHex};
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 TEXT PRIMARY KEY REFERENCES utxo_funding(new_utxo_hash) ON DELETE CASCADE,
            owner_pkh TEXT NOT NULL,
            token_id TEXT NOT NULL,
            withdrawn_in_utxo TEXT REFERENCES utxo_spending(spent_utxo_hash) ON DELETE SET NULL
        )",
        [],
    )
    .expect("failed to create table pool");

    conn.execute(
        "CREATE TABLE pool_history_entry (
            utxo TEXT PRIMARY KEY REFERENCES utxo_funding(new_utxo_hash) ON DELETE CASCADE,
            pool TEXT REFERENCES pool(creation_utxo) ON DELETE CASCADE,
            txid TEXT REFERENCES tx(txid) ON DELETE CASCADE,
            tx_pos INT NOT NULL,
            mtp_timestamp BIGINT,
            first_seen_timestamp BIGINT,
            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();
    conn.execute("CREATE INDEX idx_pool_history_entry_timestamp_sequence ON pool_history_entry(mtp_timestamp, sequence)", 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_hex()])?;
    let utxo_hex: Option<String> = row.next()?.map(|r| r.get(0).unwrap());

    match utxo_hex {
        Some(utxo) => Ok(Some(
            OutPointHash::from_hex(&utxo).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_hex(), pool_utxo.to_hex()],
    )
    .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_hex(),
            cauldron.pkh.to_hex(),
            cauldron.token_id.expect("token id for new pool missing").to_hex(),
            None::<String>,
        ]
    ).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, 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,
            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_hex(),
            pool.to_hex(),
            cauldron
                .new_utxo_txid
                .expect("txid of new pool history entry")
                .to_hex(),
            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 = "COALESCE(first_seen_timestamp, mtp_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, COALESCE(first_seen_timestamp, mtp_timestamp) as timestamp FROM pool_history_entry WHERE utxo = ?")?;
    let mut rows = stmt.query(params![utxo_hash.to_hex()])?;
    let row = rows
        .next()?
        .context("no pool history entry found for UTXO hash")?;
    let sats: u64 = row.get(1)?;
    let token_amount: u64 = row.get(2)?;
    Ok(PoolHistoryEntry {
        txid: row.get(0)?,
        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,
        COALESCE(phe.first_seen_timestamp, phe.mtp_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_hex(), start_time])?;

    let from_row = |row: &Row<'_>| -> Result<PoolHistoryEntry> {
        let sats = row.get(1)?;
        let token_amount = row.get(2)?;

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

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

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

    while let Some(row) = rows.next()? {
        history.push(from_row(row)?);
    }

    Ok(history)
}

pub fn db_pool_get_details(db: &Connection, pool: &PoolID) -> Result<(String, String)> {
    let res = db.query_row(
        "SELECT token_id, owner_pkh FROM pool WHERE creation_utxo = ?1",
        [pool.to_hex()],
        |row| Ok((row.get(0)?, row.get(1)?)),
    )?;
    Ok(res)
}

/// 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
        SUM(ABS(phe.sats_delta)) 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 COALESCE(tx.first_seen_timestamp, tx.mtp_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)> {
    let sql = "
    SELECT
        SUM(ABS(phe.sats_delta)) AS token_volume_sats,
        SUM(ABS(phe.token_delta)) AS token_volume_tokens
    FROM pool_history_entry phe
    JOIN pool p ON phe.pool = p.creation_utxo
    JOIN tx ON phe.txid = tx.txid
    WHERE COALESCE(tx.first_seen_timestamp, tx.mtp_timestamp) BETWEEN ? AND ?
        AND p.token_id = ?";

    let (sats_volume, token_volume): (i64, i64) = db.query_row(
        sql,
        params![start_timestamp, end_timestamp, token_id],
        |row| Ok((row.get(0)?, row.get(1)?)),
    )?;

    Ok((sats_volume, token_volume))
}