data_source.rs

use std::{collections::HashMap, time::Duration};

use alloy::primitives::Address;
use anyhow::Context;
use async_trait::async_trait;
use committable::Commitment;
use espresso_types::{
    Certificate2, FeeAccount, FeeAccountProof, FeeMerkleTree, Leaf2, NodeState, PubKey,
    Transaction,
    config::PublicNetworkConfig,
    v0::traits::{PersistenceOptions, SequencerPersistence},
    v0_3::{
        AuthenticatedValidator, ChainConfig, RegisteredValidator, RewardAccountProofV1,
        RewardAccountQueryDataV1, RewardAccountV1, RewardAmount, RewardMerkleTreeV1,
        StakeTableEvent,
    },
    v0_4::{RewardAccountProofV2, RewardAccountQueryDataV2, RewardAccountV2, RewardMerkleTreeV2},
};
use futures::future::{BoxFuture, Future};
use hotshot::types::BLSPubKey;
use hotshot_query_service::{
    availability::{AvailabilityDataSource, BlockQueryData, LeafQueryData, VidCommonQueryData},
    data_source::{UpdateDataSource, VersionedDataSource},
    fetching::provider::AnyProvider,
    node::NodeDataSource,
    status::StatusDataSource,
};
use hotshot_types::{
    PeerConfig,
    data::{EpochNumber, VidShare, ViewNumber},
    light_client::LCV3StateSignatureRequestBody,
    simple_certificate::LightClientStateUpdateCertificateV2,
    traits::{network::ConnectedNetwork, node_implementation::NodeType},
};
use indexmap::IndexMap;
use light_client::{state::LightClientOptions, storage::LightClientSqliteOptions};
use serde::{Deserialize, Serialize};
use tide_disco::Url;

use super::{
    AccountQueryData, BlocksFrontier, fs,
    options::{Options, Query},
    sql,
};
use crate::{
    SeqTypes, U256,
    api::{ApiState, LightClientProvider},
    persistence,
    state_cert::StateCertFetchError,
};

pub trait DataSourceOptions: PersistenceOptions {
    type DataSource: SequencerDataSource<Options = Self>;

    fn enable_query_module(&self, opt: Options, query: Query) -> Options;
}

impl DataSourceOptions for persistence::sql::Options {
    type DataSource = sql::DataSource;

    fn enable_query_module(&self, opt: Options, query: Query) -> Options {
        opt.query_sql(query, self.clone())
    }
}

impl DataSourceOptions for persistence::fs::Options {
    type DataSource = fs::DataSource;

    fn enable_query_module(&self, opt: Options, query: Query) -> Options {
        opt.query_fs(query, self.clone())
    }
}

/// A data source with sequencer-specific functionality.
///
/// This trait extends the generic [`AvailabilityDataSource`] with some additional data needed to
/// provided sequencer-specific endpoints.
#[async_trait]
pub trait SequencerDataSource:
    AvailabilityDataSource<SeqTypes>
    + NodeDataSource<SeqTypes>
    + StatusDataSource
    + UpdateDataSource<SeqTypes>
    + VersionedDataSource
    + Sized
{
    type Options: DataSourceOptions<DataSource = Self>;

    /// Instantiate a data source from command line options.
    async fn create(opt: Self::Options, provider: Provider, reset: bool) -> anyhow::Result<Self>;
}

/// Provider for fetching missing data for the query service.
pub type Provider = AnyProvider<SeqTypes>;

/// Create a provider for fetching missing data from a list of peer query services.
pub(super) async fn provider<N, P>(
    peers: impl IntoIterator<Item = Url>,
    state: &ApiState<N, P>,
    opt: LightClientOptions,
    db_opt: LightClientSqliteOptions,
) -> anyhow::Result<Provider>
where
    N: ConnectedNetwork<PubKey>,
    P: SequencerPersistence,
{
    Ok(Provider::default()
        .with_provider(LightClientProvider::new(peers, state.clone(), opt, db_opt).await?))
}

pub(crate) trait SubmitDataSource<N: ConnectedNetwork<PubKey>, P: SequencerPersistence> {
    fn submit(&self, tx: Transaction) -> impl Send + Future<Output = anyhow::Result<()>>;
}

pub(crate) trait HotShotConfigDataSource {
    fn get_config(&self) -> impl Send + Future<Output = PublicNetworkConfig>;
}

#[async_trait]
pub(crate) trait StateSignatureDataSource<N: ConnectedNetwork<PubKey>> {
    async fn get_state_signature(&self, height: u64) -> Option<LCV3StateSignatureRequestBody>;
}

pub(crate) trait NodeStateDataSource {
    fn node_state(&self) -> impl Send + Future<Output = NodeState>;
}

pub(crate) trait TokenDataSource<T: NodeType> {
    fn get_initial_supply_l1(&self) -> impl Send + Future<Output = anyhow::Result<U256>>;
    fn get_total_supply_l1(&self) -> impl Send + Future<Output = anyhow::Result<U256>>;
    fn get_decided_header(&self) -> impl Send + Future<Output = espresso_types::Header>;
}

#[derive(Serialize, Deserialize)]
#[serde(bound = "T: NodeType")]
pub struct StakeTableWithEpochNumber<T: NodeType> {
    pub epoch: Option<EpochNumber>,
    pub stake_table: Vec<PeerConfig<T>>,
}

pub(crate) trait StakeTableDataSource<T: NodeType> {
    /// Get the stake table for a given epoch
    fn get_stake_table(
        &self,
        epoch: Option<EpochNumber>,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<PeerConfig<T>>>>;

    /// Get the stake table for the current epoch if not provided
    fn get_stake_table_current(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<StakeTableWithEpochNumber<T>>>;

    /// Get the DA stake table for a given epoch
    fn get_da_stake_table(
        &self,
        epoch: Option<EpochNumber>,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<PeerConfig<T>>>>;

    /// Get the DA stake table for the current epoch if not provided
    fn get_da_stake_table_current(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<StakeTableWithEpochNumber<T>>>;

    /// Get all the validators
    fn get_validators(
        &self,
        epoch: EpochNumber,
    ) -> impl Send + Future<Output = anyhow::Result<IndexMap<Address, AuthenticatedValidator<BLSPubKey>>>>;

    fn get_block_reward(
        &self,
        epoch: Option<EpochNumber>,
    ) -> impl Send + Future<Output = anyhow::Result<Option<RewardAmount>>>;
    /// Get the current proposal participation.
    fn current_proposal_participation(
        &self,
    ) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>>;

    /// Get the proposal participation for a given epoch.
    fn proposal_participation(
        &self,
        epoch: EpochNumber,
    ) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>>;
    /// Get the current vote participation.
    fn current_vote_participation(&self) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>>;

    /// Get the vote participation for a given epoch.
    fn vote_participation(
        &self,
        epoch: EpochNumber,
    ) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>>;

    fn get_all_validators(
        &self,
        epoch: EpochNumber,
        offset: u64,
        limit: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<RegisteredValidator<PubKey>>>>;

    /// Get stake table events from L1 blocks `from_l1_block..=to_l1_block`.
    fn stake_table_events(
        &self,
        from_l1_block: u64,
        to_l1_block: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<StakeTableEvent>>>;
}

// Thin wrapper trait to access persistence methods from API handlers
#[async_trait]
pub(crate) trait StateCertDataSource {
    async fn get_state_cert_by_epoch(
        &self,
        epoch: u64,
    ) -> anyhow::Result<Option<LightClientStateUpdateCertificateV2<SeqTypes>>>;

    async fn insert_state_cert(
        &self,
        epoch: u64,
        cert: LightClientStateUpdateCertificateV2<SeqTypes>,
    ) -> anyhow::Result<()>;
}

pub(crate) trait CatchupDataSource: Sync {
    /// Get the state of the requested `account`.
    ///
    /// The state is fetched from a snapshot at the given height and view, which _must_ correspond!
    /// `height` is provided to simplify lookups for backends where data is not indexed by view.
    /// This function is intended to be used for catchup, so `view` should be no older than the last
    /// decided view.
    fn get_account(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        account: FeeAccount,
    ) -> impl Send + Future<Output = anyhow::Result<AccountQueryData>> {
        async move {
            let tree = self
                .get_accounts(instance, height, view, &[account])
                .await?;
            let (proof, balance) = FeeAccountProof::prove(&tree, account.into()).context(
                format!("account {account} not available for height {height}, view {view}"),
            )?;
            Ok(AccountQueryData { balance, proof })
        }
    }

    /// Get the state of the requested `accounts`.
    ///
    /// The state is fetched from a snapshot at the given height and view, which _must_ correspond!
    /// `height` is provided to simplify lookups for backends where data is not indexed by view.
    /// This function is intended to be used for catchup, so `view` should be no older than the last
    /// decided view.
    fn get_accounts(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[FeeAccount],
    ) -> impl Send + Future<Output = anyhow::Result<FeeMerkleTree>>;

    /// Get the blocks Merkle tree frontier.
    ///
    /// The state is fetched from a snapshot at the given height and view, which _must_ correspond!
    /// `height` is provided to simplify lookups for backends where data is not indexed by view.
    /// This function is intended to be used for catchup, so `view` should be no older than the last
    /// decided view.
    fn get_frontier(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
    ) -> impl Send + Future<Output = anyhow::Result<BlocksFrontier>>;

    fn get_chain_config(
        &self,
        commitment: Commitment<ChainConfig>,
    ) -> impl Send + Future<Output = anyhow::Result<ChainConfig>>;

    fn get_leaf_chain(
        &self,
        height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<Leaf2>>>;

    /// Load the earliest cert2 whose finalized block height is at or above `height`.
    ///
    /// Returns `None` when no cert2 height >= `height` is locally available
    fn get_cert2(
        &self,
        _height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Option<Certificate2<SeqTypes>>>> {
        async { Ok(None) }
    }

    /// Get the state of the requested `account`.
    ///
    /// The state is fetched from a snapshot at the given height and view, which _must_ correspond!
    /// `height` is provided to simplify lookups for backends where data is not indexed by view.
    /// This function is intended to be used for catchup, so `view` should be no older than the last
    /// decided view.
    fn get_reward_account_v2(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        account: RewardAccountV2,
    ) -> impl Send + Future<Output = anyhow::Result<RewardAccountQueryDataV2>> {
        async move {
            let tree = self
                .get_reward_accounts_v2(instance, height, view, &[account])
                .await?;
            let (proof, balance) = RewardAccountProofV2::prove(&tree, account.into()).context(
                format!("reward account {account} not available for height {height}, view {view}"),
            )?;
            Ok(RewardAccountQueryDataV2 { balance, proof })
        }
    }

    fn get_reward_accounts_v2(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV2],
    ) -> impl Send + Future<Output = anyhow::Result<RewardMerkleTreeV2>>;

    fn get_reward_account_v1(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        account: RewardAccountV1,
    ) -> impl Send + Future<Output = anyhow::Result<RewardAccountQueryDataV1>> {
        async move {
            let tree = self
                .get_reward_accounts_v1(instance, height, view, &[account])
                .await?;
            let (proof, balance) = RewardAccountProofV1::prove(&tree, account.into()).context(
                format!("reward account {account} not available for height {height}, view {view}"),
            )?;
            Ok(RewardAccountQueryDataV1 { balance, proof })
        }
    }

    fn get_reward_accounts_v1(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV1],
    ) -> impl Send + Future<Output = anyhow::Result<RewardMerkleTreeV1>>;

    fn get_reward_merkle_tree_v2(
        &self,
        height: u64,
        view: ViewNumber,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>>;

    fn get_state_cert(
        &self,
        epoch: u64,
    ) -> impl Send + Future<Output = anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>>>;
}

pub trait RequestResponseDataSource<Types: NodeType> {
    fn request_vid_shares(
        &self,
        block_number: u64,
        vid_common_data: VidCommonQueryData<Types>,
        duration: Duration,
    ) -> impl Future<Output = BoxFuture<'static, anyhow::Result<Vec<VidShare>>>> + Send;
}

#[async_trait]
pub trait StateCertFetchingDataSource<Types: NodeType> {
    async fn request_state_cert(
        &self,
        epoch: u64,
        timeout: Duration,
    ) -> Result<LightClientStateUpdateCertificateV2<Types>, StateCertFetchError>;
}

/// Database table size information.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct TableSize {
    pub table_name: String,
    pub row_count: i64,
    pub total_size_bytes: Option<i64>,
}

/// Status of a single deferred background migration.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct MigrationStatus {
    pub name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub completed_at: Option<chrono::DateTime<chrono::Utc>>,
    pub last_offset: Option<i64>,
}

/// Data source for database metadata and statistics.
///
/// This trait is only implemented by SQL-based storage backends (PostgreSQL and SQLite).
pub(crate) trait DatabaseMetadataSource {
    /// Get the sizes of all tables in the database.
    fn get_table_sizes(&self) -> impl Send + Future<Output = anyhow::Result<Vec<TableSize>>>;

    /// Get the status of all deferred background migrations.
    fn get_migration_status(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<MigrationStatus>>>;
}

// ============================================================================
// Arc delegation implementations
// ============================================================================
// These implementations allow Arc<T> to implement the data source traits
// when T implements them, which is necessary for NodeApiStateImpl to work
// with Arc-wrapped data sources.

use std::sync::Arc;

#[async_trait]
impl<D> StateCertDataSource for Arc<D>
where
    D: StateCertDataSource + Sync + Send,
{
    async fn get_state_cert_by_epoch(
        &self,
        epoch: u64,
    ) -> anyhow::Result<Option<LightClientStateUpdateCertificateV2<SeqTypes>>> {
        (*self).get_state_cert_by_epoch(epoch).await
    }

    async fn insert_state_cert(
        &self,
        epoch: u64,
        cert: LightClientStateUpdateCertificateV2<SeqTypes>,
    ) -> anyhow::Result<()> {
        (*self).insert_state_cert(epoch, cert).await
    }
}

impl<Types, D> RequestResponseDataSource<Types> for Arc<D>
where
    Types: NodeType,
    D: RequestResponseDataSource<Types> + Send + Sync,
{
    async fn request_vid_shares(
        &self,
        block_number: u64,
        vid_common_data: VidCommonQueryData<Types>,
        timeout_duration: Duration,
    ) -> BoxFuture<'static, anyhow::Result<Vec<VidShare>>> {
        self.as_ref()
            .request_vid_shares(block_number, vid_common_data, timeout_duration)
            .await
    }
}

#[async_trait]
impl<Types, D> StateCertFetchingDataSource<Types> for Arc<D>
where
    Types: NodeType,
    D: StateCertFetchingDataSource<Types> + Sync + Send,
{
    async fn request_state_cert(
        &self,
        epoch: u64,
        timeout: Duration,
    ) -> Result<LightClientStateUpdateCertificateV2<Types>, StateCertFetchError> {
        (*self).request_state_cert(epoch, timeout).await
    }
}

#[async_trait]
impl<T, D> StakeTableDataSource<T> for Arc<D>
where
    T: NodeType,
    D: StakeTableDataSource<T> + Sync + Send,
{
    fn get_stake_table(
        &self,
        epoch: Option<EpochNumber>,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<PeerConfig<T>>>> {
        let this = self.clone();
        async move { (*this).get_stake_table(epoch).await }
    }

    fn get_stake_table_current(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<StakeTableWithEpochNumber<T>>> {
        let this = self.clone();
        async move { (*this).get_stake_table_current().await }
    }

    fn get_da_stake_table(
        &self,
        epoch: Option<EpochNumber>,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<PeerConfig<T>>>> {
        let this = self.clone();
        async move { (*this).get_da_stake_table(epoch).await }
    }

    fn get_da_stake_table_current(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<StakeTableWithEpochNumber<T>>> {
        let this = self.clone();
        async move { (*this).get_da_stake_table_current().await }
    }

    fn get_validators(
        &self,
        epoch: EpochNumber,
    ) -> impl Send + Future<Output = anyhow::Result<IndexMap<Address, AuthenticatedValidator<BLSPubKey>>>>
    {
        let this = self.clone();
        async move { (*this).get_validators(epoch).await }
    }

    fn get_block_reward(
        &self,
        epoch: Option<EpochNumber>,
    ) -> impl Send + Future<Output = anyhow::Result<Option<RewardAmount>>> {
        let this = self.clone();
        async move { (*this).get_block_reward(epoch).await }
    }

    fn current_proposal_participation(
        &self,
    ) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>> {
        let this = self.clone();
        async move { (*this).current_proposal_participation().await }
    }

    fn proposal_participation(
        &self,
        epoch: EpochNumber,
    ) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>> {
        let this = self.clone();
        async move { (*this).proposal_participation(epoch).await }
    }

    fn current_vote_participation(&self) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>> {
        let this = self.clone();
        async move { (*this).current_vote_participation().await }
    }

    fn vote_participation(
        &self,
        epoch: EpochNumber,
    ) -> impl Send + Future<Output = HashMap<BLSPubKey, f64>> {
        let this = self.clone();
        async move { (*this).vote_participation(epoch).await }
    }

    fn get_all_validators(
        &self,
        epoch: EpochNumber,
        offset: u64,
        limit: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<RegisteredValidator<PubKey>>>> {
        let this = self.clone();
        async move { (*this).get_all_validators(epoch, offset, limit).await }
    }

    fn stake_table_events(
        &self,
        from_l1_block: u64,
        to_l1_block: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<StakeTableEvent>>> {
        let this = self.clone();
        async move { (*this).stake_table_events(from_l1_block, to_l1_block).await }
    }
}

/// Data source for pruning state: the oldest retained block and leaf.
///
/// SQL backends return the actual oldest entry; the filesystem backend always returns `None`
/// since it does not prune.
pub(crate) trait PruningDataSource {
    /// Get the oldest block in storage, or `None` if empty or unsupported.
    fn get_oldest_block(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<Option<BlockQueryData<SeqTypes>>>>;

    /// Get the oldest leaf in storage, or `None` if empty or unsupported.
    fn get_oldest_leaf(
        &self,
    ) -> impl Send + Future<Output = anyhow::Result<Option<LeafQueryData<SeqTypes>>>>;
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::{super::Options, *};

    #[async_trait]
    pub trait TestableSequencerDataSource: SequencerDataSource {
        type Storage: Sync;

        async fn create_storage() -> Self::Storage;
        fn persistence_options(storage: &Self::Storage) -> Self::Options;
        fn leaf_only_ds_options(
            _storage: &Self::Storage,
            _opt: Options,
        ) -> anyhow::Result<Options> {
            anyhow::bail!("not supported")
        }
        fn options(storage: &Self::Storage, opt: Options) -> Options;
    }
}