// Copyright The Hush Developers  2019-2022
// Released under the GPLv3

// Bounded-concurrency fetch + scan of full transactions at the end of a sync
// (LightClient::scan_fill_fulltxs). One shared gRPC channel on the caller's Tokio runtime
// replaces a fresh runtime + TLS connection per txid; parse + scan stays on the sync ThreadPool.
// Privacy: which items are requested, in what order, how often they are retried and when the
// batch stops never depend on whether an item is one of the wallet's own txs. Only the final
// Ok/Err does.

use std::collections::{BTreeSet, VecDeque};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;

use log::{info, warn};
use threadpool::ThreadPool;
use tokio::runtime::Runtime;
use zcash_primitives::transaction::{Transaction, TxId};

use crate::lightwallet::LightWallet;

/// One GetTransaction attempt; resolves to the raw transaction bytes.
pub type FetchFuture = Pin<Box<dyn Future<Output = Result<Vec<u8>, String>> + Send>>;
/// fetch(txid, reconnect): reconnect == true means "dial a new shared connection first".
pub type FetchFn = Arc<dyn Fn(TxId, bool) -> FetchFuture + Send + Sync>;
/// scan(raw_tx_bytes, height): parse + LightWallet::scan_full_tx. Always runs on the ThreadPool.
pub type ScanFn = Arc<dyn Fn(Vec<u8>, i32) -> Result<(), String> + Send + Sync>;

#[derive(Clone, Copy, Debug)]
pub struct FetchPolicy {
    pub workers: usize,
    pub max_attempts: u32,
    pub backoff_ms: u64,
    pub attempt_timeout: Duration,
    pub fail_breaker: usize,
}

impl FetchPolicy {
    /// Production policy. Fetch concurrency == the sync ThreadPool size (max(2, min(8, cpus))),
    /// i.e. exactly as many concurrent GetTransaction calls as the old one-job-per-txid code.
    pub fn for_pool(pool: &ThreadPool) -> FetchPolicy {
        FetchPolicy {
            workers: pool.max_count(),
            max_attempts: 3,
            backoff_ms: 250,
            attempt_timeout: Duration::from_secs(120),
            fail_breaker: 16,
        }
    }
}

/// Exactly the old list: own ∪ decoys, sorted by (txid, height), exact duplicates removed.
/// Returns (list, own_set); an item is "own" iff its (txid, height) is in own_set.
pub fn plan_fetch_list(own: Vec<(TxId, i32)>, decoys: Vec<(TxId, i32)>)
    -> (Vec<(TxId, i32)>, BTreeSet<(TxId, i32)>)
{
    let own_set: BTreeSet<(TxId, i32)> = own.iter().cloned().collect();
    let mut all = own;
    all.extend_from_slice(&decoys[..]);
    all.sort();
    all.dedup();
    (all, own_set)
}

/// The production scan step: identical wallet effect to the old pool job
/// (`scan_full_tx(&tx, height, 0)`), but unparsable bytes are an Err instead of a panic.
pub fn wallet_scan_fn(wallet: Arc<RwLock<LightWallet>>) -> ScanFn {
    Arc::new(move |bytes: Vec<u8>, height: i32| {
        let tx = match Transaction::read(&bytes[..]) {
            Ok(tx) => tx,
            Err(e) => return Err(format!("Error reading transaction: {:?}", e)),
        };
        wallet.read().unwrap().scan_full_tx(&tx, height, 0);
        Ok(())
    })
}

/// Fetch + scan every item. Ok iff every own item was fetched and scanned.
/// Must be called from a thread that is not inside a Tokio runtime (the sync thread).
pub fn run_fetch_scan(
    rt: &mut Runtime,
    items: Vec<(TxId, i32)>,
    own_set: BTreeSet<(TxId, i32)>,
    policy: FetchPolicy,
    pool: ThreadPool,
    fetch: FetchFn,
    scan: ScanFn,
) -> Result<(), String> {
    let own_total = items.iter().filter(|item| own_set.contains(*item)).count();
    let total = items.len();
    let workers = std::cmp::max(1, std::cmp::min(policy.workers, total));
    let queue = Arc::new(Mutex::new(items.into_iter().collect::<VecDeque<(TxId, i32)>>()));
    let own_set = Arc::new(own_set);
    let own_ok = Arc::new(AtomicUsize::new(0));
    let streak = Arc::new(AtomicUsize::new(0));
    let abort = Arc::new(AtomicBool::new(false));
    let first_err: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));

    rt.block_on(async {
        let mut handles = Vec::with_capacity(workers);
        for _ in 0..workers {
            let queue = queue.clone();
            let own_set = own_set.clone();
            let own_ok = own_ok.clone();
            let streak = streak.clone();
            let abort = abort.clone();
            let first_err = first_err.clone();
            let pool = pool.clone();
            let fetch = fetch.clone();
            let scan = scan.clone();
            handles.push(tokio::spawn(async move {
                loop {
                    if abort.load(Ordering::SeqCst) {
                        break;
                    }
                    let next = queue.lock().unwrap().pop_front();
                    let (txid, height) = match next {
                        Some(item) => item,
                        None => break,
                    };
                    let mut result: Result<(), String> = Err("not attempted".to_string());
                    for attempt in 0..policy.max_attempts {
                        if attempt > 0 {
                            tokio::time::delay_for(Duration::from_millis(policy.backoff_ms << (attempt - 1))).await;
                        }
                        let fetched = match tokio::time::timeout(policy.attempt_timeout, fetch(txid, attempt > 0)).await {
                            Ok(r) => r,
                            Err(_) => Err(format!("Error in get_transaction runtime timed out after {:?}", policy.attempt_timeout)),
                        };
                        let bytes = match fetched {
                            Ok(bytes) => bytes,
                            Err(e) => {
                                result = Err(e);
                                continue;
                            }
                        };
                        // Parse + scan on the pool. A panic there drops `done`, which turns into
                        // Err here instead of a completion that never arrives (the old wedge).
                        let (done, finished) = tokio::sync::oneshot::channel::<Result<(), String>>();
                        let scan = scan.clone();
                        pool.execute(move || {
                            let _ = done.send(scan(bytes, height));
                        });
                        result = match finished.await {
                            Ok(r) => r,
                            Err(_) => Err("full tx scan job panicked".to_string()),
                        };
                        break; // parse/scan failures are deterministic: never re-fetch
                    }
                    match result {
                        Ok(()) => {
                            streak.store(0, Ordering::SeqCst);
                            if own_set.contains(&(txid, height)) {
                                own_ok.fetch_add(1, Ordering::SeqCst);
                            }
                        }
                        Err(e) => {
                            warn!("Could not fetch full tx {}: {}", txid, e);
                            if own_set.contains(&(txid, height)) {
                                let mut first = first_err.lock().unwrap();
                                if first.is_none() {
                                    *first = Some(e);
                                }
                            }
                            // Ownership-blind breaker: a dead network stops the batch; a failed own
                            // tx alone never does (stopping on it would tell the server it is ours).
                            if streak.fetch_add(1, Ordering::SeqCst) + 1 >= policy.fail_breaker {
                                abort.store(true, Ordering::SeqCst);
                            }
                        }
                    }
                }
            }));
        }
        for handle in handles {
            let _ = handle.await; // a panicked worker shows up below as own_ok < own_total
        }
    });

    let ok = own_ok.load(Ordering::SeqCst);
    if ok < own_total {
        let first = match first_err.lock() {
            Ok(mut guard) => guard.take(),
            Err(poisoned) => poisoned.into_inner().take(),
        };
        return Err(first.unwrap_or_else(|| format!("fetched only {} of {} wallet txs", ok, own_total)));
    }
    info!("Full tx fetch done: {} items, {} new", total, own_total);
    Ok(())
}