rostrum 14.0.1

use std::{collections::HashMap, sync::Arc};

use crate::{
    chaindef::{OutPointHash, ScriptHash, TokenID},
    query::{queryutil::outpoints_are_spent, status::scan_for_outputs, BUFFER_SIZE, CHUNK_SIZE},
    store::{DBContents, DBStore},
    util::Channel,
};

use super::{
    queryfilter::QueryFilter,
    tip::{get_utxos_in_db, UtxoEntry},
};
use anyhow::Result;
use futures::{stream, StreamExt};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use tokio::try_join;

/// When no height filter is used, we can use the faster unspent output index.
pub async fn balance_at_tip(
    confirmed: &Arc<DBStore>,
    mempool: &Arc<DBStore>,
    scripthash: &ScriptHash,
    filter: &QueryFilter,
    filter_token: &Option<TokenID>,
) -> Result<(u64, HashMap<TokenID, u64>, i64, HashMap<TokenID, i64>)> {
    assert!(!filter.has_height_filter());
    let c_channel: Channel<Option<UtxoEntry>> = Channel::bounded(1000);
    let c_sender = c_channel.sender();

    // Find current confirmed balance
    let confirmed_task = tokio::spawn(async move {
        let mut c_receiver = c_channel.receiver().await;

        let mut bch: u64 = 0;
        let mut tokens: HashMap<TokenID, u64> = HashMap::default();
        #[allow(clippy::type_complexity)]
        let mut outputs: Vec<(OutPointHash, (u64, Option<TokenID>, Option<i64>))> = Vec::default();
        while let Some(Some((unspent, utxo, token))) = c_receiver.recv().await {
            bch += utxo.value();
            if let Some((token_id, amount, _commitment)) = token.as_ref() {
                // nexa authority utxo is recognized via negative amount
                let amount = if *amount < 0 { 0 } else { *amount as u64 };
                if let Some(b) = tokens.get_mut(token_id) {
                    *b += amount
                } else {
                    tokens.insert(token_id.clone(), amount);
                }
            }
            let token_amount = token.as_ref().map(|t| t.1);
            let token_id = token.map(|t| t.0);
            outputs.push((
                unspent.outpointhash(),
                (utxo.value(), token_id, token_amount),
            ))
        }
        (bch, tokens, outputs)
    });

    // Find current unconfirmed balance
    let u_channel: Channel<Option<UtxoEntry>> = Channel::bounded(1000);
    let u_sender = u_channel.sender();
    let unconfirmed_task = tokio::spawn(async move {
        let mut u_receiver = u_channel.receiver().await;

        let mut bch: i64 = 0;
        let mut tokens: HashMap<TokenID, i64> = HashMap::default();

        while let Some(Some((_unspent, utxo, token))) = u_receiver.recv().await {
            bch += utxo.value() as i64;
            if let Some((token_id, amount, _commitment)) = token {
                // nexa authority utxo is recognized via negative amount
                let amount = if amount < 0 { 0 } else { amount };
                if let Some(b) = tokens.get_mut(&token_id) {
                    *b += amount;
                } else {
                    tokens.insert(token_id, amount);
                }
            }
        }
        (bch, tokens)
    });
    try_join!(
        get_utxos_in_db(
            mempool,
            None,
            scripthash,
            filter,
            filter_token,
            u_sender.clone(),
            true
        ),
        get_utxos_in_db(
            confirmed,
            None,
            scripthash,
            filter,
            filter_token,
            c_sender.clone(),
            true
        )
    )?;
    let (
        (confirmed_bch, confirmed_token, confirmed_outpoints),
        (mut unconfirmed_bch, mut unconfirmed_token),
    ) = try_join!(confirmed_task, unconfirmed_task)?;

    // Subtract any spends of confirmed in mempool (balance can be negative)
    let (outpoints, amounts): (Vec<OutPointHash>, Vec<_>) = confirmed_outpoints.into_iter().unzip();
    let spent_status = outpoints_are_spent(mempool, &outpoints).await;
    for (is_spent, (bch, token, token_amount)) in spent_status.into_iter().zip(amounts.into_iter())
    {
        if !is_spent {
            continue;
        }
        unconfirmed_bch -= bch as i64;
        if let Some(amount) = token_amount {
            let token_id = token.unwrap();
            if let Some(v) = unconfirmed_token.get_mut(&token_id) {
                *v -= amount
            } else {
                unconfirmed_token.insert(token_id, -amount);
            }
        }
    }

    Ok((
        confirmed_bch,
        confirmed_token,
        unconfirmed_bch,
        unconfirmed_token,
    ))
}

/**
 * Calculate coin balance in scripthash that has been confirmed in blocks.
 *
 * This call uses full spent/unspent indexes to query historical balances.
 */
pub async fn historical_confirmed_balance(
    index: &Arc<DBStore>,
    scripthash: ScriptHash,
    filter: &QueryFilter,
) -> Result<(i64, Vec<(i64, OutPointHash)>)> {
    assert!(index.contents == DBContents::ConfirmedIndex);

    let (outputs_query, outputs_stream) = scan_for_outputs(index, scripthash, filter).await;

    let outputs: Vec<(i64, OutPointHash)> = outputs_stream
        .map(|(o, _)| (o.value(), o.take_hash()))
        .ready_chunks(CHUNK_SIZE)
        .then(|outpoints| {
            // Check if outpoints are spent (in baches of CHUNK_SIZE)
            let index = Arc::clone(index);
            async move {
                let (values, outpoints): (Vec<u64>, Vec<OutPointHash>) =
                    outpoints.into_iter().unzip();
                let spent_status = outpoints_are_spent(&index, &outpoints).await;

                stream::iter(values.into_iter().zip(outpoints).zip(spent_status).map(
                    |((value, outpoint), is_spent)| {
                        if is_spent {
                            // Output was created AND spent in the same index. Zero it out.
                            (0, outpoint)
                        } else {
                            (value as i64, outpoint)
                        }
                    },
                ))
            }
        })
        .flatten_unordered(BUFFER_SIZE)
        .collect()
        .await;

    outputs_query.await??;

    let amount = outputs.par_iter().map(|(value, _)| value).sum();

    // Only return unspent confirmed outputs
    let outputs = outputs
        .into_iter()
        .filter(|(amount, _)| amount > &0)
        .collect();

    Ok((amount, outputs))
}

/**
 * Calculate coin balance in scripthash that has been confirmed in blocks.
 *
 * Takes confirmed_outputs in to be able to see if mempool spends confirmed utxos.
 *
 * The 'historical' is because it's needed for queries that filter on QueryFilter::from_height
 */
pub async fn historical_unconfirmed_balance(
    mempool: &Arc<DBStore>,
    confirmed_outputs: Vec<(i64, OutPointHash)>,
    scripthash: ScriptHash,
    filter: &QueryFilter,
) -> Result<i64> {
    assert!(DBContents::MempoolIndex == mempool.contents);

    let (outputs_query, outputs_stream) = scan_for_outputs(mempool, scripthash, filter).await;

    let unconfirmed_outputs: Vec<i64> = outputs_stream
        .map(|(o, _)| (o.value(), o.take_hash()))
        .ready_chunks(CHUNK_SIZE)
        .then(|outpoints| {
            let mempool = mempool.clone();
            async move {
                let (values, outpoints): (Vec<u64>, Vec<OutPointHash>) =
                    outpoints.into_iter().unzip();
                let is_spent = outpoints_are_spent(&mempool, &outpoints).await;
                stream::iter(values.into_iter().zip(is_spent).map(|(value, is_spent)| {
                    if is_spent {
                        // Output was created AND spent in the same index. Zero it out.
                        0
                    } else {
                        value as i64
                    }
                }))
            }
        })
        .flatten_unordered(BUFFER_SIZE)
        .collect()
        .await;

    outputs_query.await??;

    let amount: i64 = unconfirmed_outputs.par_iter().sum();

    // Subtract spends from confirmed utxos
    let spent_confirmed: i64 = stream::iter(confirmed_outputs)
        .ready_chunks(CHUNK_SIZE)
        .then(|outpoints| {
            let mempool = mempool.clone();
            async move {
                let (value, outpoints): (Vec<i64>, Vec<OutPointHash>) =
                    outpoints.into_iter().unzip();
                let is_spent = outpoints_are_spent(&mempool, &outpoints).await;

                stream::iter(value.into_iter().zip(is_spent).map(|(value, is_spent)| {
                    debug_assert!(value >= 0);
                    if is_spent {
                        value
                    } else {
                        0
                    }
                }))
            }
        })
        .flatten_unordered(BUFFER_SIZE)
        .fold(0i64, |acc, value| async move { acc + value })
        .await;

    Ok(amount - spent_confirmed)
}