#![deny(clippy::pedantic)]
#![allow(clippy::cast_possible_wrap)]
#![allow(clippy::module_name_repetitions)]

pub use fedimint_lnv2_common as common;

mod db;

use std::collections::{BTreeMap, BTreeSet};
use std::time::Duration;

use anyhow::{Context, anyhow, ensure};
use bls12_381::{G1Projective, Scalar};
use fedimint_core::bitcoin::hashes::sha256;
use fedimint_core::config::{
    ServerModuleConfig, ServerModuleConsensusConfig, TypedServerModuleConfig,
    TypedServerModuleConsensusConfig,
};
use fedimint_core::core::ModuleInstanceId;
use fedimint_core::db::{
    Database, DatabaseTransaction, DatabaseVersion, IDatabaseTransactionOpsCoreTyped,
};
use fedimint_core::encoding::Encodable;
use fedimint_core::envs::{FM_ENABLE_MODULE_LNV2_ENV, is_env_var_set_opt};
use fedimint_core::module::audit::Audit;
use fedimint_core::module::{
    Amounts, ApiEndpoint, ApiError, ApiVersion, CORE_CONSENSUS_VERSION, CoreConsensusVersion,
    InputMeta, ModuleConsensusVersion, ModuleInit, SupportedModuleApiVersions,
    TransactionItemAmounts, api_endpoint,
};
use fedimint_core::net::auth::check_auth;
use fedimint_core::task::timeout;
use fedimint_core::time::duration_since_epoch;
use fedimint_core::util::SafeUrl;
use fedimint_core::{
    BitcoinHash, InPoint, NumPeers, NumPeersExt, OutPoint, PeerId, apply, async_trait_maybe_send,
    push_db_pair_items,
};
use fedimint_lnv2_common::config::{
    FeeConsensus, LightningClientConfig, LightningConfig, LightningConfigConsensus,
    LightningConfigPrivate,
};
use fedimint_lnv2_common::contracts::{IncomingContract, OutgoingContract};
use fedimint_lnv2_common::endpoint_constants::{
    ADD_GATEWAY_ENDPOINT, AWAIT_INCOMING_CONTRACT_ENDPOINT, AWAIT_INCOMING_CONTRACTS_ENDPOINT,
    AWAIT_PREIMAGE_ENDPOINT, CONSENSUS_BLOCK_COUNT_ENDPOINT, DECRYPTION_KEY_SHARE_ENDPOINT,
    GATEWAYS_ENDPOINT, OUTGOING_CONTRACT_EXPIRATION_ENDPOINT, REMOVE_GATEWAY_ENDPOINT,
};
use fedimint_lnv2_common::{
    ContractId, LightningCommonInit, LightningConsensusItem, LightningInput, LightningInputError,
    LightningInputV0, LightningModuleTypes, LightningOutput, LightningOutputError,
    LightningOutputOutcome, LightningOutputV0, MODULE_CONSENSUS_VERSION, OutgoingWitness,
};
use fedimint_logging::LOG_MODULE_LNV2;
use fedimint_server_core::bitcoin_rpc::ServerBitcoinRpcMonitor;
use fedimint_server_core::config::{PeerHandleOps, eval_poly_g1};
use fedimint_server_core::migration::ServerModuleDbMigrationFn;
use fedimint_server_core::{
    ConfigGenModuleArgs, ServerModule, ServerModuleInit, ServerModuleInitArgs,
};
use futures::StreamExt;
use group::Curve;
use group::ff::Field;
use rand::SeedableRng;
use rand_chacha::ChaChaRng;
use strum::IntoEnumIterator;
use tpe::{
    AggregatePublicKey, DecryptionKeyShare, PublicKeyShare, SecretKeyShare, derive_pk_share,
};
use tracing::trace;

use crate::db::{
    BlockCountVoteKey, BlockCountVotePrefix, DbKeyPrefix, DecryptionKeyShareKey,
    DecryptionKeySharePrefix, GatewayKey, GatewayPrefix, IncomingContractIndexKey,
    IncomingContractIndexPrefix, IncomingContractKey, IncomingContractOutpointKey,
    IncomingContractOutpointPrefix, IncomingContractPrefix, IncomingContractStreamIndexKey,
    IncomingContractStreamKey, IncomingContractStreamPrefix, OutgoingContractKey,
    OutgoingContractPrefix, PreimageKey, PreimagePrefix, UnixTimeVoteKey, UnixTimeVotePrefix,
};

#[derive(Debug, Clone)]
pub struct LightningInit;

impl ModuleInit for LightningInit {
    type Common = LightningCommonInit;

    #[allow(clippy::too_many_lines)]
    async fn dump_database(
        &self,
        dbtx: &mut DatabaseTransaction<'_>,
        prefix_names: Vec<String>,
    ) -> Box<dyn Iterator<Item = (String, Box<dyn erased_serde::Serialize + Send>)> + '_> {
        let mut lightning: BTreeMap<String, Box<dyn erased_serde::Serialize + Send>> =
            BTreeMap::new();

        let filtered_prefixes = DbKeyPrefix::iter().filter(|f| {
            prefix_names.is_empty() || prefix_names.contains(&f.to_string().to_lowercase())
        });

        for table in filtered_prefixes {
            match table {
                DbKeyPrefix::BlockCountVote => {
                    push_db_pair_items!(
                        dbtx,
                        BlockCountVotePrefix,
                        BlockCountVoteKey,
                        u64,
                        lightning,
                        "Lightning Block Count Votes"
                    );
                }
                DbKeyPrefix::UnixTimeVote => {
                    push_db_pair_items!(
                        dbtx,
                        UnixTimeVotePrefix,
                        UnixTimeVoteKey,
                        u64,
                        lightning,
                        "Lightning Unix Time Votes"
                    );
                }
                DbKeyPrefix::OutgoingContract => {
                    push_db_pair_items!(
                        dbtx,
                        OutgoingContractPrefix,
                        LightningOutgoingContractKey,
                        OutgoingContract,
                        lightning,
                        "Lightning Outgoing Contracts"
                    );
                }
                DbKeyPrefix::IncomingContract => {
                    push_db_pair_items!(
                        dbtx,
                        IncomingContractPrefix,
                        LightningIncomingContractKey,
                        IncomingContract,
                        lightning,
                        "Lightning Incoming Contracts"
                    );
                }
                DbKeyPrefix::IncomingContractOutpoint => {
                    push_db_pair_items!(
                        dbtx,
                        IncomingContractOutpointPrefix,
                        LightningIncomingContractOutpointKey,
                        OutPoint,
                        lightning,
                        "Lightning Incoming Contracts Outpoints"
                    );
                }
                DbKeyPrefix::DecryptionKeyShare => {
                    push_db_pair_items!(
                        dbtx,
                        DecryptionKeySharePrefix,
                        DecryptionKeyShareKey,
                        DecryptionKeyShare,
                        lightning,
                        "Lightning Decryption Key Share"
                    );
                }
                DbKeyPrefix::Preimage => {
                    push_db_pair_items!(
                        dbtx,
                        PreimagePrefix,
                        LightningPreimageKey,
                        [u8; 32],
                        lightning,
                        "Lightning Preimages"
                    );
                }
                DbKeyPrefix::Gateway => {
                    push_db_pair_items!(
                        dbtx,
                        GatewayPrefix,
                        GatewayKey,
                        (),
                        lightning,
                        "Lightning Gateways"
                    );
                }
                DbKeyPrefix::IncomingContractStreamIndex => {
                    push_db_pair_items!(
                        dbtx,
                        IncomingContractStreamIndexKey,
                        IncomingContractStreamIndexKey,
                        u64,
                        lightning,
                        "Lightning Incoming Contract Stream Index"
                    );
                }
                DbKeyPrefix::IncomingContractStream => {
                    push_db_pair_items!(
                        dbtx,
                        IncomingContractStreamPrefix(0),
                        IncomingContractStreamKey,
                        IncomingContract,
                        lightning,
                        "Lightning Incoming Contract Stream"
                    );
                }
                DbKeyPrefix::IncomingContractIndex => {
                    push_db_pair_items!(
                        dbtx,
                        IncomingContractIndexPrefix,
                        IncomingContractIndexKey,
                        u64,
                        lightning,
                        "Lightning Incoming Contract Index"
                    );
                }
            }
        }

        Box::new(lightning.into_iter())
    }
}

#[apply(async_trait_maybe_send!)]
impl ServerModuleInit for LightningInit {
    type Module = Lightning;

    fn versions(&self, _core: CoreConsensusVersion) -> &[ModuleConsensusVersion] {
        &[MODULE_CONSENSUS_VERSION]
    }

    fn supported_api_versions(&self) -> SupportedModuleApiVersions {
        SupportedModuleApiVersions::from_raw(
            (CORE_CONSENSUS_VERSION.major, CORE_CONSENSUS_VERSION.minor),
            (
                MODULE_CONSENSUS_VERSION.major,
                MODULE_CONSENSUS_VERSION.minor,
            ),
            &[(0, 0)],
        )
    }

    fn is_enabled_by_default(&self) -> bool {
        is_env_var_set_opt(FM_ENABLE_MODULE_LNV2_ENV).unwrap_or(true)
    }

    async fn init(&self, args: &ServerModuleInitArgs<Self>) -> anyhow::Result<Self::Module> {
        Ok(Lightning {
            cfg: args.cfg().to_typed()?,
            db: args.db().clone(),
            server_bitcoin_rpc_monitor: args.server_bitcoin_rpc_monitor(),
        })
    }

    fn trusted_dealer_gen(
        &self,
        peers: &[PeerId],
        args: &ConfigGenModuleArgs,
    ) -> BTreeMap<PeerId, ServerModuleConfig> {
        let tpe_pks = peers
            .iter()
            .map(|peer| (*peer, dealer_pk(peers.to_num_peers(), *peer)))
            .collect::<BTreeMap<PeerId, PublicKeyShare>>();

        peers
            .iter()
            .map(|peer| {
                let cfg = LightningConfig {
                    consensus: LightningConfigConsensus {
                        tpe_agg_pk: dealer_agg_pk(),
                        tpe_pks: tpe_pks.clone(),
                        fee_consensus: if args.disable_base_fees {
                            FeeConsensus::zero()
                        } else {
                            FeeConsensus::new(0).expect("Relative fee is within range")
                        },
                        network: args.network,
                    },
                    private: LightningConfigPrivate {
                        sk: dealer_sk(peers.to_num_peers(), *peer),
                    },
                };

                (*peer, cfg.to_erased())
            })
            .collect()
    }

    async fn distributed_gen(
        &self,
        peers: &(dyn PeerHandleOps + Send + Sync),
        args: &ConfigGenModuleArgs,
    ) -> anyhow::Result<ServerModuleConfig> {
        let (polynomial, sks) = peers.run_dkg_g1().await?;

        let server = LightningConfig {
            consensus: LightningConfigConsensus {
                tpe_agg_pk: tpe::AggregatePublicKey(polynomial[0].to_affine()),
                tpe_pks: peers
                    .num_peers()
                    .peer_ids()
                    .map(|peer| (peer, PublicKeyShare(eval_poly_g1(&polynomial, &peer))))
                    .collect(),
                fee_consensus: if args.disable_base_fees {
                    FeeConsensus::zero()
                } else {
                    FeeConsensus::new(0).expect("Relative fee is within range")
                },
                network: args.network,
            },
            private: LightningConfigPrivate {
                sk: SecretKeyShare(sks),
            },
        };

        Ok(server.to_erased())
    }

    fn validate_config(&self, identity: &PeerId, config: ServerModuleConfig) -> anyhow::Result<()> {
        let config = config.to_typed::<LightningConfig>()?;

        ensure!(
            tpe::derive_pk_share(&config.private.sk)
                == *config
                    .consensus
                    .tpe_pks
                    .get(identity)
                    .context("Public key set has no key for our identity")?,
            "Preimge encryption secret key share does not match our public key share"
        );

        Ok(())
    }

    fn get_client_config(
        &self,
        config: &ServerModuleConsensusConfig,
    ) -> anyhow::Result<LightningClientConfig> {
        let config = LightningConfigConsensus::from_erased(config)?;
        Ok(LightningClientConfig {
            tpe_agg_pk: config.tpe_agg_pk,
            tpe_pks: config.tpe_pks,
            fee_consensus: config.fee_consensus,
            network: config.network,
        })
    }

    fn get_database_migrations(
        &self,
    ) -> BTreeMap<DatabaseVersion, ServerModuleDbMigrationFn<Lightning>> {
        let mut migrations: BTreeMap<DatabaseVersion, ServerModuleDbMigrationFn<Lightning>> =
            BTreeMap::new();

        migrations.insert(
            DatabaseVersion(0),
            Box::new(move |ctx| Box::pin(crate::db::migrate_to_v1(ctx))),
        );

        migrations
    }

    fn used_db_prefixes(&self) -> Option<BTreeSet<u8>> {
        Some(DbKeyPrefix::iter().map(|p| p as u8).collect())
    }
}

fn dealer_agg_pk() -> AggregatePublicKey {
    AggregatePublicKey((G1Projective::generator() * coefficient(0)).to_affine())
}

fn dealer_pk(num_peers: NumPeers, peer: PeerId) -> PublicKeyShare {
    derive_pk_share(&dealer_sk(num_peers, peer))
}

fn dealer_sk(num_peers: NumPeers, peer: PeerId) -> SecretKeyShare {
    let x = Scalar::from(peer.to_usize() as u64 + 1);

    // We evaluate the scalar polynomial of degree threshold - 1 at the point x
    // using the Horner schema.

    let y = (0..num_peers.threshold())
        .map(|index| coefficient(index as u64))
        .rev()
        .reduce(|accumulator, c| accumulator * x + c)
        .expect("We have at least one coefficient");

    SecretKeyShare(y)
}

fn coefficient(index: u64) -> Scalar {
    Scalar::random(&mut ChaChaRng::from_seed(
        *index.consensus_hash::<sha256::Hash>().as_byte_array(),
    ))
}

#[derive(Debug)]
pub struct Lightning {
    cfg: LightningConfig,
    db: Database,
    server_bitcoin_rpc_monitor: ServerBitcoinRpcMonitor,
}

#[apply(async_trait_maybe_send!)]
impl ServerModule for Lightning {
    type Common = LightningModuleTypes;
    type Init = LightningInit;

    async fn consensus_proposal(
        &self,
        _dbtx: &mut DatabaseTransaction<'_>,
    ) -> Vec<LightningConsensusItem> {
        // We reduce the time granularity to deduplicate votes more often and not save
        // one consensus item every second.
        let mut items = vec![LightningConsensusItem::UnixTimeVote(
            60 * (duration_since_epoch().as_secs() / 60),
        )];

        if let Ok(block_count) = self.get_block_count() {
            trace!(target: LOG_MODULE_LNV2, ?block_count, "Proposing block count");
            items.push(LightningConsensusItem::BlockCountVote(block_count));
        }

        items
    }

    async fn process_consensus_item<'a, 'b>(
        &'a self,
        dbtx: &mut DatabaseTransaction<'b>,
        consensus_item: LightningConsensusItem,
        peer: PeerId,
    ) -> anyhow::Result<()> {
        trace!(target: LOG_MODULE_LNV2, ?consensus_item, "Processing consensus item proposal");

        match consensus_item {
            LightningConsensusItem::BlockCountVote(vote) => {
                let current_vote = dbtx
                    .insert_entry(&BlockCountVoteKey(peer), &vote)
                    .await
                    .unwrap_or(0);

                ensure!(current_vote < vote, "Block count vote is redundant");

                Ok(())
            }
            LightningConsensusItem::UnixTimeVote(vote) => {
                let current_vote = dbtx
                    .insert_entry(&UnixTimeVoteKey(peer), &vote)
                    .await
                    .unwrap_or(0);

                ensure!(current_vote < vote, "Unix time vote is redundant");

                Ok(())
            }
            LightningConsensusItem::Default { variant, .. } => Err(anyhow!(
                "Received lnv2 consensus item with unknown variant {variant}"
            )),
        }
    }

    async fn process_input<'a, 'b, 'c>(
        &'a self,
        dbtx: &mut DatabaseTransaction<'c>,
        input: &'b LightningInput,
        _in_point: InPoint,
    ) -> Result<InputMeta, LightningInputError> {
        let (pub_key, amount) = match input.ensure_v0_ref()? {
            LightningInputV0::Outgoing(outpoint, outgoing_witness) => {
                let contract = dbtx
                    .remove_entry(&OutgoingContractKey(*outpoint))
                    .await
                    .ok_or(LightningInputError::UnknownContract)?;

                let pub_key = match outgoing_witness {
                    OutgoingWitness::Claim(preimage) => {
                        if contract.expiration <= self.consensus_block_count(dbtx).await {
                            return Err(LightningInputError::Expired);
                        }

                        if !contract.verify_preimage(preimage) {
                            return Err(LightningInputError::InvalidPreimage);
                        }

                        dbtx.insert_entry(&PreimageKey(*outpoint), preimage).await;

                        contract.claim_pk
                    }
                    OutgoingWitness::Refund => {
                        if contract.expiration > self.consensus_block_count(dbtx).await {
                            return Err(LightningInputError::NotExpired);
                        }

                        contract.refund_pk
                    }
                    OutgoingWitness::Cancel(forfeit_signature) => {
                        if !contract.verify_forfeit_signature(forfeit_signature) {
                            return Err(LightningInputError::InvalidForfeitSignature);
                        }

                        contract.refund_pk
                    }
                };

                (pub_key, contract.amount)
            }
            LightningInputV0::Incoming(outpoint, agg_decryption_key) => {
                let contract = dbtx
                    .remove_entry(&IncomingContractKey(*outpoint))
                    .await
                    .ok_or(LightningInputError::UnknownContract)?;

                let index = dbtx
                    .remove_entry(&IncomingContractIndexKey(*outpoint))
                    .await
                    .expect("Incoming contract index should exist");

                dbtx.remove_entry(&IncomingContractStreamKey(index)).await;

                if !contract
                    .verify_agg_decryption_key(&self.cfg.consensus.tpe_agg_pk, agg_decryption_key)
                {
                    return Err(LightningInputError::InvalidDecryptionKey);
                }

                let pub_key = match contract.decrypt_preimage(agg_decryption_key) {
                    Some(..) => contract.commitment.claim_pk,
                    None => contract.commitment.refund_pk,
                };

                (pub_key, contract.commitment.amount)
            }
        };

        Ok(InputMeta {
            amount: TransactionItemAmounts {
                amounts: Amounts::new_bitcoin(amount),
                fees: Amounts::new_bitcoin(self.cfg.consensus.fee_consensus.fee(amount)),
            },
            pub_key,
        })
    }

    async fn process_output<'a, 'b>(
        &'a self,
        dbtx: &mut DatabaseTransaction<'b>,
        output: &'a LightningOutput,
        outpoint: OutPoint,
    ) -> Result<TransactionItemAmounts, LightningOutputError> {
        let amount = match output.ensure_v0_ref()? {
            LightningOutputV0::Outgoing(contract) => {
                dbtx.insert_new_entry(&OutgoingContractKey(outpoint), contract)
                    .await;

                contract.amount
            }
            LightningOutputV0::Incoming(contract) => {
                if !contract.verify() {
                    return Err(LightningOutputError::InvalidContract);
                }

                if contract.commitment.expiration <= self.consensus_unix_time(dbtx).await {
                    return Err(LightningOutputError::ContractExpired);
                }

                dbtx.insert_new_entry(&IncomingContractKey(outpoint), contract)
                    .await;

                dbtx.insert_entry(
                    &IncomingContractOutpointKey(contract.contract_id()),
                    &outpoint,
                )
                .await;

                let stream_index = dbtx
                    .get_value(&IncomingContractStreamIndexKey)
                    .await
                    .unwrap_or(0);

                dbtx.insert_entry(&IncomingContractStreamKey(stream_index), contract)
                    .await;

                dbtx.insert_entry(&IncomingContractIndexKey(outpoint), &stream_index)
                    .await;

                dbtx.insert_entry(&IncomingContractStreamIndexKey, &(stream_index + 1))
                    .await;

                let dk_share = contract.create_decryption_key_share(&self.cfg.private.sk);

                dbtx.insert_entry(&DecryptionKeyShareKey(outpoint), &dk_share)
                    .await;

                contract.commitment.amount
            }
        };

        Ok(TransactionItemAmounts {
            amounts: Amounts::new_bitcoin(amount),
            fees: Amounts::new_bitcoin(self.cfg.consensus.fee_consensus.fee(amount)),
        })
    }

    async fn output_status(
        &self,
        _dbtx: &mut DatabaseTransaction<'_>,
        _out_point: OutPoint,
    ) -> Option<LightningOutputOutcome> {
        None
    }

    async fn audit(
        &self,
        dbtx: &mut DatabaseTransaction<'_>,
        audit: &mut Audit,
        module_instance_id: ModuleInstanceId,
    ) {
        // Both incoming and outgoing contracts represent liabilities to the federation
        // since they are obligations to issue notes.
        audit
            .add_items(
                dbtx,
                module_instance_id,
                &OutgoingContractPrefix,
                |_, contract| -(contract.amount.msats as i64),
            )
            .await;

        audit
            .add_items(
                dbtx,
                module_instance_id,
                &IncomingContractPrefix,
                |_, contract| -(contract.commitment.amount.msats as i64),
            )
            .await;
    }

    fn api_endpoints(&self) -> Vec<ApiEndpoint<Self>> {
        vec![
            api_endpoint! {
                CONSENSUS_BLOCK_COUNT_ENDPOINT,
                ApiVersion::new(0, 0),
                async |module: &Lightning, context, _params : () | -> u64 {
                    let db = context.db();
                    let mut dbtx = db.begin_transaction_nc().await;

                    Ok(module.consensus_block_count(&mut dbtx).await)
                }
            },
            api_endpoint! {
                AWAIT_INCOMING_CONTRACT_ENDPOINT,
                ApiVersion::new(0, 0),
                async |module: &Lightning, context, params: (ContractId, u64) | -> Option<OutPoint> {
                    let db = context.db();

                    Ok(module.await_incoming_contract(db, params.0, params.1).await)
                }
            },
            api_endpoint! {
                AWAIT_PREIMAGE_ENDPOINT,
                ApiVersion::new(0, 0),
                async |module: &Lightning, context, params: (OutPoint, u64)| -> Option<[u8; 32]> {
                    let db = context.db();

                    Ok(module.await_preimage(db, params.0, params.1).await)
                }
            },
            api_endpoint! {
                DECRYPTION_KEY_SHARE_ENDPOINT,
                ApiVersion::new(0, 0),
                async |_module: &Lightning, context, params: OutPoint| -> DecryptionKeyShare {
                    let share = context
                        .db()
                        .begin_transaction_nc()
                        .await
                        .get_value(&DecryptionKeyShareKey(params))
                        .await
                        .ok_or(ApiError::bad_request("No decryption key share found".to_string()))?;

                    Ok(share)
                }
            },
            api_endpoint! {
                OUTGOING_CONTRACT_EXPIRATION_ENDPOINT,
                ApiVersion::new(0, 0),
                async |module: &Lightning, context, outpoint: OutPoint| -> Option<(ContractId, u64)> {
                    let db = context.db();

                    Ok(module.outgoing_contract_expiration(db, outpoint).await)
                }
            },
            api_endpoint! {
                AWAIT_INCOMING_CONTRACTS_ENDPOINT,
                ApiVersion::new(0, 0),
                async |module: &Lightning, context, params: (u64, usize)| -> (Vec<IncomingContract>, u64) {
                    let db = context.db();

                    if params.1 == 0 {
                        return Err(ApiError::bad_request("Batch size must be greater than 0".to_string()));
                    }

                    Ok(module.await_incoming_contracts(db, params.0, params.1).await)
                }
            },
            api_endpoint! {
                ADD_GATEWAY_ENDPOINT,
                ApiVersion::new(0, 0),
                async |_module: &Lightning, context, gateway: SafeUrl| -> bool {
                    check_auth(context)?;

                    let db = context.db();

                    Ok(Lightning::add_gateway(db, gateway).await)
                }
            },
            api_endpoint! {
                REMOVE_GATEWAY_ENDPOINT,
                ApiVersion::new(0, 0),
                async |_module: &Lightning, context, gateway: SafeUrl| -> bool {
                    check_auth(context)?;

                    let db = context.db();

                    Ok(Lightning::remove_gateway(db, gateway).await)
                }
            },
            api_endpoint! {
                GATEWAYS_ENDPOINT,
                ApiVersion::new(0, 0),
                async |_module: &Lightning, context, _params : () | -> Vec<SafeUrl> {
                    let db = context.db();

                    Ok(Lightning::gateways(db).await)
                }
            },
        ]
    }
}

impl Lightning {
    fn get_block_count(&self) -> anyhow::Result<u64> {
        self.server_bitcoin_rpc_monitor
            .status()
            .map(|status| status.block_count)
            .context("Block count not available yet")
    }

    async fn consensus_block_count(&self, dbtx: &mut DatabaseTransaction<'_>) -> u64 {
        let num_peers = self.cfg.consensus.tpe_pks.to_num_peers();

        let mut counts = dbtx
            .find_by_prefix(&BlockCountVotePrefix)
            .await
            .map(|entry| entry.1)
            .collect::<Vec<u64>>()
            .await;

        counts.sort_unstable();

        counts.reverse();

        assert!(counts.last() <= counts.first());

        // The block count we select guarantees that any threshold of correct peers can
        // increase the consensus block count and any consensus block count has been
        // confirmed by a threshold of peers.

        counts.get(num_peers.threshold() - 1).copied().unwrap_or(0)
    }

    async fn consensus_unix_time(&self, dbtx: &mut DatabaseTransaction<'_>) -> u64 {
        let num_peers = self.cfg.consensus.tpe_pks.to_num_peers();

        let mut times = dbtx
            .find_by_prefix(&UnixTimeVotePrefix)
            .await
            .map(|entry| entry.1)
            .collect::<Vec<u64>>()
            .await;

        times.sort_unstable();

        times.reverse();

        assert!(times.last() <= times.first());

        // The unix time we select guarantees that any threshold of correct peers can
        // advance the consensus unix time and any consensus unix time has been
        // confirmed by a threshold of peers.

        times.get(num_peers.threshold() - 1).copied().unwrap_or(0)
    }

    async fn await_incoming_contract(
        &self,
        db: Database,
        contract_id: ContractId,
        expiration: u64,
    ) -> Option<OutPoint> {
        loop {
            timeout(
                Duration::from_secs(10),
                db.wait_key_exists(&IncomingContractOutpointKey(contract_id)),
            )
            .await
            .ok();

            // to avoid race conditions we have to check for the contract and
            // its expiration in the same database transaction
            let mut dbtx = db.begin_transaction_nc().await;

            if let Some(outpoint) = dbtx
                .get_value(&IncomingContractOutpointKey(contract_id))
                .await
            {
                return Some(outpoint);
            }

            if expiration <= self.consensus_unix_time(&mut dbtx).await {
                return None;
            }
        }
    }

    async fn await_preimage(
        &self,
        db: Database,
        outpoint: OutPoint,
        expiration: u64,
    ) -> Option<[u8; 32]> {
        loop {
            timeout(
                Duration::from_secs(10),
                db.wait_key_exists(&PreimageKey(outpoint)),
            )
            .await
            .ok();

            // to avoid race conditions we have to check for the preimage and
            // the contracts expiration in the same database transaction
            let mut dbtx = db.begin_transaction_nc().await;

            if let Some(preimage) = dbtx.get_value(&PreimageKey(outpoint)).await {
                return Some(preimage);
            }

            if expiration <= self.consensus_block_count(&mut dbtx).await {
                return None;
            }
        }
    }

    async fn outgoing_contract_expiration(
        &self,
        db: Database,
        outpoint: OutPoint,
    ) -> Option<(ContractId, u64)> {
        let mut dbtx = db.begin_transaction_nc().await;

        let contract = dbtx.get_value(&OutgoingContractKey(outpoint)).await?;

        let consensus_block_count = self.consensus_block_count(&mut dbtx).await;

        let expiration = contract.expiration.saturating_sub(consensus_block_count);

        Some((contract.contract_id(), expiration))
    }

    async fn await_incoming_contracts(
        &self,
        db: Database,
        start: u64,
        n: usize,
    ) -> (Vec<IncomingContract>, u64) {
        let filter = |next_index: Option<u64>| next_index.filter(|i| *i > start);

        let (mut next_index, mut dbtx) = db
            .wait_key_check(&IncomingContractStreamIndexKey, filter)
            .await;

        let mut contracts = Vec::with_capacity(n);

        let range = IncomingContractStreamKey(start)..IncomingContractStreamKey(u64::MAX);

        for (key, contract) in dbtx
            .find_by_range(range)
            .await
            .take(n)
            .collect::<Vec<(IncomingContractStreamKey, IncomingContract)>>()
            .await
        {
            contracts.push(contract.clone());
            next_index = key.0 + 1;
        }

        (contracts, next_index)
    }

    async fn add_gateway(db: Database, gateway: SafeUrl) -> bool {
        let mut dbtx = db.begin_transaction().await;

        let is_new_entry = dbtx.insert_entry(&GatewayKey(gateway), &()).await.is_none();

        dbtx.commit_tx().await;

        is_new_entry
    }

    async fn remove_gateway(db: Database, gateway: SafeUrl) -> bool {
        let mut dbtx = db.begin_transaction().await;

        let entry_existed = dbtx.remove_entry(&GatewayKey(gateway)).await.is_some();

        dbtx.commit_tx().await;

        entry_existed
    }

    async fn gateways(db: Database) -> Vec<SafeUrl> {
        db.begin_transaction_nc()
            .await
            .find_by_prefix(&GatewayPrefix)
            .await
            .map(|entry| entry.0.0)
            .collect()
            .await
    }

    pub async fn consensus_block_count_ui(&self) -> u64 {
        self.consensus_block_count(&mut self.db.begin_transaction_nc().await)
            .await
    }

    pub async fn consensus_unix_time_ui(&self) -> u64 {
        self.consensus_unix_time(&mut self.db.begin_transaction_nc().await)
            .await
    }

    pub async fn add_gateway_ui(&self, gateway: SafeUrl) -> bool {
        Self::add_gateway(self.db.clone(), gateway).await
    }

    pub async fn remove_gateway_ui(&self, gateway: SafeUrl) -> bool {
        Self::remove_gateway(self.db.clone(), gateway).await
    }

    pub async fn gateways_ui(&self) -> Vec<SafeUrl> {
        Self::gateways(self.db.clone()).await
    }
}
