rostrum 14.0.1

use crate::{
    discovery::{client::ElectrumClient, utilparser::lookup_hosts},
    metrics::Metrics,
    signal::Waiter,
};
use anyhow::Result;
use bitcoincash::Network;
use rand::RngExt;
use rayon::prelude::IntoParallelRefIterator;
use serde_json::Value;
use std::{
    collections::{HashMap, HashSet},
    convert::TryInto,
    net::IpAddr,
    sync::{atomic::AtomicBool, Arc},
    time::{Duration, Instant},
};
use tokio::sync::Mutex;

use crate::config::Config;
use crate::discovery::{
    peer::{CandidatePeer, GoodKnownPeer, ServerPeer},
    seed::{SEED_PEERS_CHIPNET, SEED_PEERS_MAINNET, SEED_PEERS_TESTNET, SEED_PEERS_TESTNET4},
};
use rayon::prelude::*;

// We only accept limited number of peers per subnet as sybil attack mitigation
const MAX_PER_OCTET_PREFIX: usize = 5;
// Max number of candidates in queue for consideration
const MAX_CANDIDATES: usize = 100;
// If candidate queue is full, accept a new one only if this many seconds have passed
const CANDIDATE_EVICT_DURATION: Duration = Duration::from_secs(120);
// How long until we re-visit a good peer to reconsider if it's good or not.
const GOOD_PEER_REVISIT_DURATION: Duration = Duration::from_secs(3600);
// How many servers can be queued for our announcement
const ANNOUNCE_QUEUE_MAX_SIZE: usize = 200;

// Get IP-prefix
fn get_prefix(addr: &IpAddr) -> [u8; 2] {
    match addr {
        IpAddr::V4(ipv4) => ipv4.octets()[..2].try_into().unwrap(),
        IpAddr::V6(ipv6) => ipv6.octets()[..2].try_into().unwrap(),
    }
}

/**
 * Maintains a list of good server peers and considers new ones.
 */
pub struct PeerDiscoverer {
    good: Mutex<HashMap<u64, GoodKnownPeer>>,
    /// List of peers that have submitted themselves to us for consideration.
    candidates: Mutex<Vec<CandidatePeer>>,
    /// Last time we added someone to the candidate list.
    last_add: Mutex<Instant>,
    /// Thread responsible for visiting server peers
    discovery_thread: Mutex<Option<tokio::task::JoinHandle<()>>>,
    discovery_thread_kill: Arc<AtomicBool>,
    /// Yet to announce to
    announce_queue: Mutex<Vec<CandidatePeer>>,
    /// Peers we have (attempted to) announce to.
    announce_done: Mutex<HashSet</* peer id */ u64>>,

    /// How many good nodes we currently know of.
    metric_good_now: prometheus::IntGauge,
    /// How many good nodes we have added in total.
    metric_good_added: prometheus::IntCounter,
    /// How many good nodes we removed total.
    metric_good_removed: prometheus::IntCounter,
    /// How many candidates we currently have.
    metric_candidates_now: prometheus::IntGauge,
    /// How many candidates we have added in total.
    metric_candidates_added: prometheus::IntCounter,
    /// How many candidates we have rejected.
    metric_candidates_rejected: prometheus::IntCounter,
    /// How many peers are queued for announcement.
    metric_announce_now: prometheus::IntGauge,
}

// TODO: Store good peers in a database
impl PeerDiscoverer {
    pub async fn new(config: &Config, metrics: &Metrics) -> Arc<Self> {
        let d = Arc::new(Self {
            good: Mutex::new(HashMap::default()),
            candidates: Mutex::new(match config.network_type {
                Network::Bitcoin => SEED_PEERS_MAINNET.to_vec(),
                Network::Testnet => SEED_PEERS_TESTNET.to_vec(),
                Network::Testnet4 => SEED_PEERS_TESTNET4.to_vec(),
                Network::Chipnet => SEED_PEERS_CHIPNET.to_vec(),
                Network::Regtest | Network::Scalenet => Vec::default(),
            }),
            last_add: Mutex::new(Instant::now()),
            discovery_thread: Mutex::new(None),
            discovery_thread_kill: Arc::new(AtomicBool::new(false)),
            announce_done: Mutex::new(HashSet::default()),
            announce_queue: Mutex::new(Vec::with_capacity(ANNOUNCE_QUEUE_MAX_SIZE)),

            metric_good_now: metrics.gauge_int(prometheus::Opts::new(
                "rostrum_discover_good",
                "# of known good server peers",
            )),
            metric_good_added: metrics.counter_int(prometheus::Opts::new(
                "rostrum_discover_added",
                "# of known good server peers added (accumulative)",
            )),
            metric_good_removed: metrics.counter_int(prometheus::Opts::new(
                "rostrum_discover_removed",
                "# of known good server peers removed (accumulative)",
            )),
            metric_candidates_now: metrics.gauge_int(prometheus::Opts::new(
                "rostrum_discover_candidates",
                "# of known untested server peers",
            )),
            metric_candidates_added: metrics.counter_int(prometheus::Opts::new(
                "rostrum_discover_candidates_added",
                "# of server peer candidates added (accumulative)",
            )),
            metric_candidates_rejected: metrics.counter_int(prometheus::Opts::new(
                "rostrum_discover_candidates_rejected",
                "# of server peer candidates rejected (accumulative)",
            )),
            metric_announce_now: metrics.gauge_int(prometheus::Opts::new(
                "rostrum_discover_announce_queue",
                "# of server peers in queue for announcement",
            )),
        });

        let metric_tested = metrics.counter_int(prometheus::Opts::new(
            "rostrum_discover_candidates_tested",
            "# of server peer test connections made (accumulative)",
        ));

        let loop_interval = Duration::from_secs(if config.network_type == Network::Regtest {
            1
        } else {
            10
        });

        let thread_d = Arc::clone(&d);
        if config.announce {
            info!("Peer discovery thread started.");
            *d.discovery_thread.lock().await =
                Some(crate::thread::spawn_task("discovery", async move {
                    loop {
                        if Waiter::wait_or_shutdown(loop_interval).await.is_err()
                            || thread_d
                                .discovery_thread_kill
                                .load(std::sync::atomic::Ordering::Relaxed)
                        {
                            // Exit thread
                            return Ok(());
                        }

                        // Checking if good nodes are still good is takes precedence over finding a new candidate.
                        let needs_revisit = thread_d
                            .good
                            .lock()
                            .await
                            .iter()
                            .filter_map(|(_, p)| {
                                if p.last_visit.elapsed() >= GOOD_PEER_REVISIT_DURATION {
                                    // Time to re-visit this peer server
                                    Some(p.clone())
                                } else {
                                    None
                                }
                            })
                            .take(1)
                            .collect::<Vec<GoodKnownPeer>>();

                        if let Some(peer) = needs_revisit.first() {
                            metric_tested.inc();
                            if let Some(good_peer) = evaluate_peer(peer).await {
                                if good_peer.id() != peer.id() {
                                    // peer is good; but it's info has changed
                                    thread_d.remove_good(&peer.id()).await;
                                    thread_d.promote_candidate(0, good_peer).await;
                                } else {
                                    thread_d.bump_last_visit(&peer.id()).await;
                                }
                            } else {
                                thread_d.remove_good(&peer.id()).await
                            }
                            continue;
                        }

                        // Good nodes are up-to-date. Check a random candidate.
                        let candidate = {
                            let lock = thread_d.candidates.lock().await;
                            if lock.is_empty() {
                                None
                            } else {
                                // We take a copy so that it isn't re-added while we're considering it
                                let index = rand::rng().random_range(0..lock.len());
                                lock.get(index).cloned()
                            }
                        };

                        if let Some(peer) = candidate {
                            if let Some(good_peer) = evaluate_peer(&peer).await {
                                thread_d.promote_candidate(peer.id(), good_peer).await;
                                let announce = thread_d.maybe_add_announce_queue(peer).await;
                                trace!("discover: Announce ourself to promoted peer: {}", announce);
                            } else {
                                thread_d.reject_candidate(&peer.id()).await;
                            }
                        }
                    }
                }));
        } else {
            info!("Peer discovery thread disabled.")
        }

        d
    }

    /// Update the 'last visit' attribute of a good node
    pub async fn bump_last_visit(&self, id: &u64) {
        match self.good.lock().await.get_mut(id) {
            Some(p) => p.last_visit = Instant::now(),
            None => {
                debug_assert!(false);
                warn!("Failed to bump good peer {id}");
            }
        };
    }

    /// Promote a candidate peer into good node
    pub async fn promote_candidate(&self, candidate_id: u64, good_peer: GoodKnownPeer) {
        if candidate_id != 0 {
            let mut lock = self.candidates.lock().await;
            let candidate = lock
                .iter()
                .enumerate()
                .find(|(_, p)| p.id() == candidate_id);
            if let Some((i, _)) = candidate {
                lock.remove(i);
            } else {
                debug_assert!(false);
                warn!("Tried to promote a candidate that does not exist");
                return;
            }
        };

        self.good.lock().await.insert(good_peer.id(), good_peer);
        self.metric_good_added.inc();
        self.metric_good_now.inc();
    }

    /// Reject a candidate
    pub async fn reject_candidate(&self, id: &u64) {
        let mut lock = self.candidates.lock().await;
        let candidate = lock.iter().enumerate().find(|(_, p)| &p.id() == id);
        match candidate {
            Some((index, _)) => {
                lock.remove(index);
                self.metric_candidates_rejected.inc();
                self.metric_candidates_now.set(lock.len() as i64);
            }
            None => {
                debug_assert!(false);
                warn!("Tried to reject a candidate that does not exist")
            }
        }
    }

    /// Remove a node no longer considered good
    pub async fn remove_good(&self, id: &u64) {
        self.good.lock().await.remove(id);
        self.metric_good_now.dec();
        self.metric_good_removed.inc();

        // If it went offline, allow us to announce ourselves to it if we see it again.
        self.announce_done.lock().await.remove(id);
    }

    /// How many good peers we have with given prefix.
    async fn len_good_peers_with_prefix(&self, prefix: &[u8; 2]) -> usize {
        self.good
            .lock()
            .await
            .par_iter()
            .filter(|(_, p)| &get_prefix(&p.ip()) == prefix)
            .count()
    }

    /// If we have a good peer with given ID
    async fn has_good(&self, id: u64) -> bool {
        self.good.lock().await.iter().any(|(_, p)| p.id() == id)
    }

    /// Add a new peer for discoverer to consider
    pub async fn add_candidate(&self, candidate: CandidatePeer) -> Result<()> {
        if self
            .len_good_peers_with_prefix(&get_prefix(&candidate.ip()))
            .await
            >= MAX_PER_OCTET_PREFIX
        {
            bail!("Have enough good servers in this subnet");
        }

        let candidate_id = candidate.id();

        if self.has_good(candidate_id).await {
            return Ok(());
        }

        let mut c = self.candidates.lock().await;
        if c.iter().any(|existing| existing.id() == candidate_id) {
            return Ok(());
        }

        if c.len() < MAX_CANDIDATES {
            c.push(candidate);
            *self.last_add.lock().await = Instant::now();
            self.metric_candidates_added.inc();
            self.metric_candidates_now.inc();
            return Ok(());
        }

        // Our candidate list is full, but if it's been a long time since we
        // accepted a candidate, just randomly drop some other and accept this one.
        if self.last_add.lock().await.elapsed() >= CANDIDATE_EVICT_DURATION {
            trace!("Evicting random and adding candidate {:?}", candidate);
            let index = rand::rng().random_range(0..c.len());
            c.remove(index);
            c.push(candidate);
            *self.last_add.lock().await = Instant::now();
            self.metric_candidates_added.inc();
            return Ok(());
        }
        bail!("Too many peers for consideration. Try again later.");
    }

    // Get list of good peers.
    pub async fn get_peers(&self) -> Vec<GoodKnownPeer> {
        self.good.lock().await.values().cloned().collect()
    }

    /// Take a peer from the announce queue.
    pub async fn next_announce(&self) -> Option<CandidatePeer> {
        let mut lock = self.announce_queue.lock().await;

        if lock.is_empty() {
            return None;
        }
        // take a random peer to mitigate sybil
        let index = rand::rng().random_range(0..lock.len());
        let peer = lock.remove(index);

        self.announce_done.lock().await.insert(peer.id());
        self.metric_announce_now.set(lock.len() as i64);
        Some(peer)
    }

    /// Add peer to queue for peers to announce this server to (if applicable).
    pub async fn maybe_add_announce_queue(&self, peer: CandidatePeer) -> bool {
        if self.announce_done.lock().await.contains(&peer.id()) {
            return false;
        }
        let mut lock = self.announce_queue.lock().await;
        if lock.len() >= ANNOUNCE_QUEUE_MAX_SIZE {
            trace!("discover: Announce queue full, ignoring {}", peer.ip());
            return false;
        }
        lock.push(peer);
        self.metric_announce_now.set(lock.len() as i64);
        true
    }

    /// Take a copy of a random good peer
    pub async fn random_good(&self) -> Option<GoodKnownPeer> {
        let lock = self.good.lock().await;
        if lock.is_empty() {
            return None;
        }
        let index = rand::rng().random_range(0..lock.len());
        lock.iter()
            .skip(index)
            .take(1)
            .map(|(_, p)| p)
            .collect::<Vec<_>>()
            .first()
            .cloned()
            .cloned()
    }
}

impl Drop for PeerDiscoverer {
    fn drop(&mut self) {
        self.discovery_thread_kill
            .store(true, std::sync::atomic::Ordering::Relaxed);
    }
}

async fn try_fetch_features(
    conn_result: Result<ElectrumClient>,
    peername: &str,
    label: &str,
) -> Option<Value> {
    match conn_result {
        Ok(conn) => match conn.call("server.features", vec![]).await {
            Ok(features) => Some(features),
            Err(e) => {
                trace!("discover: {label} on {peername} failed features call: {e}");
                None
            }
        },
        Err(e) => {
            trace!("discover: {label} on {peername} failed to connect: {e}");
            None
        }
    }
}

pub async fn evalute_peer_inner(
    peer: &dyn ServerPeer,
    allow_new_candidate: bool,
) -> (Option<GoodKnownPeer>, Option<CandidatePeer>) {
    let peername = peer.pretty_hostname();
    trace!("discover: Evaluating peer {peername}");

    let (ssl_result, tcp_result) = ElectrumClient::connect_all(peer).await;

    let features_ssl = try_fetch_features(ssl_result, &peername, "ssl").await;
    let features_tcp = try_fetch_features(tcp_result, &peername, "tcp").await;

    // If both work; but response is mismatch, discard the TCP
    let (features_ssl, features_tcp) = match (&features_ssl, &features_tcp) {
        (Some(ssl), Some(tcp)) if ssl != tcp => {
            trace!(
                "discover: mismatch between ssl and tcp on {}",
                peer.pretty_hostname()
            );
            (features_ssl.clone(), None)
        }
        _ => (features_ssl, features_tcp),
    };

    let features = if let Some(features) = features_ssl.clone().or(features_tcp.clone()) {
        features
    } else {
        trace!("discover: Did not get features from tcp or ssl from {peername}");
        return (None, None);
    };

    let peer_version = match features
        .get("protocol_max")
        .and_then(|protocol| protocol.as_str())
    {
        Some(version) => version,
        None => {
            trace!("discover: peer {peername} did not provide version");
            return (None, None);
        }
    };

    let hosts = lookup_hosts(features.get("hosts"), peer.ip()).await;

    // prefer hostname
    let preferred = hosts.iter().find(|c| c.hostname.is_some());
    let preferred = preferred.unwrap_or(hosts.first().unwrap());

    // did 'features' provide different information than what we had? (which might be gossip)
    // we may want to evaluate that instead
    if allow_new_candidate && (preferred.id() != peer.id()) {
        trace!("discover: Learned new info about {peername}; recursing");
        return (None, Some(preferred.clone()));
    }

    // all looks good; let's promote
    // (if features variables are not set, means the port was bad)
    let tcp_port = features_tcp.as_ref().and_then(|_| peer.tcp_port());
    let ssl_port = features_ssl.as_ref().and_then(|_| peer.ssl_port());
    let good = GoodKnownPeer {
        tcp_port,
        ssl_port,
        hostname: preferred.hostname(),
        ip: peer.ip(),
        last_visit: Instant::now(),
        last_candidate_request: peer.last_candidate_request(),
        version: peer_version.to_string(),
    };
    (Some(good), None)
}

pub async fn evaluate_peer(peer: &dyn ServerPeer) -> Option<GoodKnownPeer> {
    match evalute_peer_inner(peer, true).await {
        (None, None) => None,
        (None, Some(candidate)) => evalute_peer_inner(&candidate, false).await.0,
        (Some(good), None) => Some(good),
        (Some(good), Some(_)) => {
            debug_assert!(false, "should not find both good and candidate");
            Some(good)
        }
    }
}