//! Reaching peers past NATs: picking relays and keeping slots on them, upgrading
//! relayed connections to direct ones, retrying hole punches, and finding peers on the
//! DHT.

use super::*;

/// How a node relates to relays.
pub(super) struct Relaying {
    /// It relays for others.
    pub(super) serving: bool,
    /// Public relays may be picked, after members.
    pub(super) public: bool,
    pub(super) local: PeerId,
}

/// The relays to pick among `offers`, if this node needs relays, keeping `using`: a
/// node needs them when it has no address others could dial (`direct`), or AutoNAT says
/// it can't be dialled; up to `AUTO_RELAYS`, in an order of its own (the lowest hash of
/// both peer IDs), so that not every node picks the same.
pub(super) fn relay_check(
    local: PeerId,
    direct: bool,
    nat: &::libp2p::autonat::NatStatus,
    offers: &HashMap<PeerId, Multiaddr>,
    using: &[PeerId],
) -> Vec<(PeerId, Multiaddr)> {
    use ::libp2p::autonat::NatStatus;
    let needed = match nat {
        NatStatus::Private => true,
        NatStatus::Public(_) => false,
        NatStatus::Unknown => !direct,
    };
    if !needed || using.len() >= AUTO_RELAYS {
        return Vec::new();
    }
    let score = |relay: &PeerId| {
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        (local, relay).hash(&mut hasher);
        hasher.finish()
    };
    let mut candidates: Vec<(PeerId, Multiaddr)> = offers
        .iter()
        .filter(|(relay, _)| **relay != local && !using.contains(relay))
        .map(|(relay, address)| (*relay, address.clone()))
        .collect();
    candidates.sort_by_key(|(relay, _)| score(relay));
    candidates.truncate(AUTO_RELAYS - using.len());
    candidates
}

/// Reserves slots on the relays `relay_check` picks (configured ones first, then listed
/// ones, members, public ones), and lets those picked go once AutoNAT says this node can
/// be dialled.
pub(super) fn pick_relays(
    swarm: &mut Swarm<Behaviour>,
    book: &Mutex<Book>,
    relaying: &Relaying,
    nat: &::libp2p::autonat::NatStatus,
    picked: &mut HashMap<::libp2p::core::transport::ListenerId, PeerId>,
) {
    if nat.is_public() && !picked.is_empty() {
        info!("Reachable directly (AutoNAT); letting the relays picked go");
        let mut book = lock(book);
        for (listener, relay) in std::mem::take(picked) {
            swarm.remove_listener(listener);
            book.slots.remove(&relay);
        }
        return;
    }
    let (direct, configured, listed, members, public) = {
        let mut book = lock(book);
        let now = tokio::time::Instant::now();
        book.failed_relays.retain(|_, until| now < *until);
        let failed = book.failed_relays.clone();
        // Not paused, and not in a range left out (ignored, or without a route): those
        // would only be denied or fail, again and again.
        let usable = |book: &mut Book, offers: HashMap<PeerId, Multiaddr>| {
            offers
                .into_iter()
                .filter(|(relay, address)| !failed.contains_key(relay) && !book.left_out(address))
                .collect::<HashMap<PeerId, Multiaddr>>()
        };
        let public = if relaying.public {
            let offers = book.public_relays.clone();
            usable(&mut book, offers)
        } else {
            HashMap::new()
        };
        let (configured, listed, members) = (
            book.named(Named::Configured),
            book.named(Named::Listed),
            book.relays.clone(),
        );
        (
            !book.direct().is_empty(),
            usable(&mut book, configured),
            usable(&mut book, listed),
            usable(&mut book, members),
            public,
        )
    };
    let mut using: Vec<PeerId> = picked.values().copied().collect();
    // Each kind for what the ones before left.
    for (kind, offers) in [
        (Named::Configured.as_str(), configured),
        (Named::Listed.as_str(), listed),
        ("member", members),
        ("public", public),
    ] {
        for (relay, address) in relay_check(relaying.local, direct, nat, &offers, &using) {
            using.push(relay);
            match swarm.listen_on(address.clone().with(Protocol::P2pCircuit)) {
                Ok(listener) => {
                    info!(relay = %Addressed(SCHEME, &relay), address = %address, kind, "Picked a relay");
                    picked.insert(listener, relay);
                }
                Err(error) => {
                    debug!(relay = %Addressed(SCHEME, &relay), kind, error = %error, "Couldn't use a relay");
                }
            }
        }
    }
}

/// Dials `peer` directly if it's connected only through a relay and is known elsewhere
/// too (the member list, identify, mDNS), since the first connection, made through a
/// relay (a relayed seed), would otherwise be kept for good: a peer connected isn't
/// dialled again, and hole punching needs both sides to listen. The direct connection,
/// once open, wins over the relayed one (`Book::duplicates`). A dial refused because
/// another to the peer is under way (hole punching's, say) doesn't count as a try: the
/// next comes when that one fails (`driver`), not `UPGRADE_RETRY` later.
pub(super) fn upgrade(
    swarm: &mut Swarm<Behaviour>,
    book: &Mutex<Book>,
    peer: PeerId,
    local: PeerId,
) {
    let (direct, now, wait) = {
        let mut book = lock(book);
        if book.connection(&peer) != Some(Connection::Relayed) {
            return;
        }
        let mut direct = book.usable(&peer, local);
        direct.retain(|address| !is_relayed(address));
        if direct.is_empty() {
            return;
        }
        let now = tokio::time::Instant::now();
        let (due, wait) = book
            .upgrades
            .get(&peer)
            .copied()
            .unwrap_or((now, UPGRADE_RETRY));
        if now < due {
            return;
        }
        (direct, now, wait)
    };
    let dialled = swarm.dial(
        DialOpts::peer_id(peer)
            .addresses(direct.clone())
            .condition(PeerCondition::NotDialing)
            .build(),
    );
    match dialled {
        Ok(()) => {
            debug!(peer = %Addressed(SCHEME, &peer), addresses = ?direct, "Connected only through a relay; dialling directly");
            lock(book)
                .upgrades
                .insert(peer, (now + wait, (wait * 2).min(RELAY_PAUSE)));
        }
        Err(error) => {
            debug!(peer = %Addressed(SCHEME, &peer), error = %error, "Connected only through a relay; can't dial directly yet");
        }
    }
}

/// Dials `peer`, connected directly, at its other addresses, in case one is a better
/// path: one that came since (a new interface of its, identify tells), or, `again`
/// (one of ours came), every one, as our side may reach them better now, the one in
/// use too. Each once while connected (`Book::explored`); a connection that opens is
/// pinged at once, and of several the fastest is kept (`Book::duplicates`). Without
/// it, a connection stays on the path it was made on (the WLAN) when a better one comes
/// (ethernet): a connected peer isn't dialled again.
pub(super) fn explore(swarm: &mut Swarm<Behaviour>, book: &Mutex<Book>, peer: PeerId, again: bool) {
    let candidates = lock(book).to_explore(peer, again);
    if candidates.is_empty() {
        return;
    }
    debug!(peer = %Addressed(SCHEME, &peer), addresses = ?candidates, "Connected directly; dialling other addresses, for a better path");
    if let Err(error) = swarm.dial(
        DialOpts::peer_id(peer)
            .addresses(candidates)
            .condition(PeerCondition::Always)
            .build(),
    ) {
        debug!(peer = %Addressed(SCHEME, &peer), error = %error, "Couldn't dial for a better path");
    }
}

/// Tries hole punching again with the peers it failed with that are due (`punching`).
/// libp2p punches as a relayed connection opens, started by the side that was dialled
/// through the relay; that side is where a punch's failure is told, so it is this one.
/// A new relayed connection starts a new punch: we dial the peer through a relay if it
/// has a slot on one; else we close our relayed connections to it, and its next message
/// (a probe, within a minute) dials us through ours again. Peers connected directly by
/// now are dropped; those not connected now are kept, failures and all.
pub(super) fn retry_punches(swarm: &mut Swarm<Behaviour>, book: &Mutex<Book>, local: PeerId) {
    let now = tokio::time::Instant::now();
    let mut close: Vec<ConnectionId> = Vec::new();
    let due: Vec<(PeerId, usize, Vec<Multiaddr>)> = {
        let mut book = lock(book);
        let book = &mut *book;
        let mut due = Vec::new();
        book.punching.retain(|peer, (failed, next)| {
            match book_connection(&book.connections, peer) {
                Some(Connection::Direct) => return false,
                // Kept, failures and all: one we closed to have it dial anew comes back.
                None => return true,
                Some(Connection::Relayed) if now < *next => return true,
                Some(Connection::Relayed) => {}
            }
            // As `Book::usable`, by parts: `punching` is borrowed.
            let through: Vec<Multiaddr> = book
                .addresses
                .get(peer)
                .map(|known| dialable(known, local))
                .unwrap_or_default()
                .into_iter()
                .filter(|address| {
                    is_relayed(address) && !excluded(&book.ignored, &mut book.routes, address)
                })
                .collect();
            // Until its punch fails again (or the peer tries from its side).
            *next = now + PUNCH_RETRIES[(*failed).min(PUNCH_RETRIES.len() - 1)];
            if through.is_empty() {
                close.extend(
                    book.connections
                        .iter()
                        .filter(|(_, (connected, relayed, _))| connected == peer && *relayed)
                        .map(|(id, _)| *id),
                );
            } else {
                due.push((*peer, *failed, through));
            }
            true
        });
        due
    };
    for connection in close {
        debug!(connection = %connection, "Hole punching again: closing a relayed connection, for the peer to dial anew");
        swarm.close_connection(connection);
    }
    for (peer, failed, through) in due {
        debug!(peer = %Addressed(SCHEME, &peer), failed, "Hole punching again, through a new relayed connection");
        let _ = swarm.dial(
            DialOpts::peer_id(peer)
                .addresses(through)
                .condition(PeerCondition::Always)
                .build(),
        );
    }
}

/// Asks the DHT (with `--libp2p-public`) where `peer` is, for the next attempt.
pub(super) fn find_on_dht(swarm: &mut Swarm<Behaviour>, peer: PeerId) {
    if let Some(dht) = swarm.behaviour_mut().kad.as_mut() {
        dht.get_closest_peers(peer);
    }
}