//! Discovering peers, the handshake (HELLO / HELLO_REPLY), and membership.

use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;

use anyhow::{Context, Result, anyhow, bail};
use tokio::sync::Mutex;
use tokio::time::Instant;
use tracing::{debug, info, instrument, trace, warn};

use crate::state::{AppState, HelloState, Liveness, PeerState, Standing};
use crate::swarm::sync::SyncEvent;
use crate::swarm::topology;
use crate::transport::{self, ParsedMessage};

/// Most peers a HELLO_REPLY lists, and most we take from one: each listed peer may get
/// a HELLO, so a long list (a hub's hundreds of neighbours, or a peer making up keys)
/// costs the recipient as much as it is long. A joining node needs only a few to find
/// the swarm; the member list gives it the rest once it is ready.
const MAX_LISTED_PEERS: usize = 64;

pub(crate) async fn remember_peer(state: &AppState, peer: &str) {
    if peer.is_empty() {
        warn!("Ignoring empty peer public key");
        return;
    }

    if peer == state.local_id() {
        debug!("Ignoring own public key");
        return;
    }

    let mut peer_states = state.peer_states.write().await;

    if !peer_states.contains_key(peer) {
        peer_states.insert(peer.to_owned(), state.new_peer_state(peer));
        info!(
            peer = %peer,
            peer_count = peer_states.len(),
            "Discovered peer"
        );
    } else {
        trace!(peer = %peer, "Peer already registered");
    }
}

/// Starts a new handshake with `peer` even though one is complete: its reply resets
/// sync on both sides and brings its current peer list (`swarm::source`, after giving
/// it up as the catch-up source).
pub(crate) async fn rehandshake(state: &AppState, peer: &str) {
    if let Some(peer_state) = state.get_peer_state(peer).await {
        let mut peer_state = peer_state.lock().await;
        if peer_state.hello_state == HelloState::Established {
            peer_state.hello_state = HelloState::NotSent;
        }
    }
    ensure_handshake(state, peer).await;
}

pub(crate) async fn send_hello_once(state: &AppState, peer: &str) -> Result<()> {
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        trace!(peer = %peer, "Peer was forgotten; not sending HELLO");
        return Ok(());
    };

    let claimed = {
        let mut peer_state = peer_state_mutex.lock().await;
        if peer_state.hello_state == HelloState::NotSent {
            // A new handshake: retries start over.
            // Numbered across all peers, so a peer we forgot and meet again still
            // sees a higher number than any it had from this run.
            peer_state.handshake = state.handshakes.fetch_add(1, Ordering::Relaxed) + 1;
            peer_state.hello_state = HelloState::Sending;
            peer_state.hello_backoff.reset();
            peer_state.hello_attempts = 0;
            true
        } else {
            trace!(peer = %peer, hello_state = ?peer_state.hello_state, "not sending");
            false
        }
    };

    if !claimed {
        return Ok(());
    }

    debug!(to = %peer, "Sending HELLO");
    transmit_hello(state, peer).await
}

pub(super) async fn retry_hello_if_due(state: &AppState, peer: &str) -> Result<()> {
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        return Ok(());
    };

    {
        let mut peer_state = peer_state_mutex.lock().await;
        match peer_state.hello_state {
            HelloState::AwaitingReply { next_attempt } if Instant::now() >= next_attempt => {
                peer_state.hello_state = HelloState::Sending;
            }
            _ => return Ok(()),
        }
    }

    debug!(peer = %peer, "No HELLO_REPLY yet; re-sending HELLO");
    transmit_hello(state, peer).await
}

/// Sends a HELLO without holding any lock, then schedules the next retry.
/// Success and failure both end in AwaitingReply: Mycelium accepting the message
/// says nothing about whether the peer received it.
async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
    let swarm_id = state.swarm_id().await;
    let auth = state
        .swarm_key
        .as_ref()
        .map(|key| key.hello_tag(swarm_id.as_deref().unwrap_or(""), state.local_id(), peer))
        .unwrap_or_default();
    let (epoch, epoch_seq) = state.epoch_stamp().await;
    // Retries repeat the handshake's number.
    let handshake = match state.get_peer_state(peer).await {
        Some(peer_state) => peer_state.lock().await.handshake,
        None => 0,
    };
    state.lock_rtt().hello_sent(peer, Instant::now());
    let result = state
        .transport
        .send_message(
            peer,
            &ParsedMessage::Hello(transport::HelloRequest {
                swarm_id,
                epoch,
                epoch_seq,
                auth,
                ready: state.is_ready().await,
                run_id: state.run_id,
                handshake,
            }),
        )
        .await;

    if let Some(peer_state_mutex) = state.get_peer_state(peer).await {
        let mut peer_state = peer_state_mutex.lock().await;

        if peer_state.hello_state == HelloState::Sending {
            let was_unreachable = peer_state.unreachable_since.is_some();
            peer_state.hello_sent(Instant::now());
            if !was_unreachable && peer_state.unreachable_since.is_some() {
                state.topology_changed.notify_one();
            }
            let wait = match peer_state.liveness {
                Liveness::Dormant { .. } => state.timing.dormant_probe,
                Liveness::Active => peer_state.hello_backoff.next(),
            };
            peer_state.hello_state = HelloState::AwaitingReply {
                next_attempt: Instant::now() + wait,
            };
        }
    }

    result.context("Sending HELLO")
}

pub(super) async fn ensure_handshake(state: &AppState, peer: &str) {
    if let Err(error) = send_hello_once(state, peer).await {
        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
    }
}

/// Answers a HELLO that carried `run_id`; `away` is how long before it we had last heard
/// from the peer.
pub(super) async fn send_hello_reply(
    state: &AppState,
    peer: &str,
    run_id: transport::RunId,
    away: Option<Duration>,
) -> Result<()> {
    let swarm_id = state
        .swarm_id()
        .await
        .ok_or_else(|| anyhow!("Cannot send HELLO_REPLY without an established swarm ID"))?;

    let candidates: Vec<(String, Arc<Mutex<PeerState>>)> = state
        .peers_snapshot()
        .await
        .into_iter()
        .filter(|(known, _)| known.as_str() != peer)
        .collect();

    let mut peers = Vec::new();
    for (known, peer_state_mutex) in candidates {
        let peer_state = peer_state_mutex.lock().await;
        if peer_state.liveness == Liveness::Active
            && peer_state.hello_state == HelloState::Established
            && topology::is_neighbour(state, &known)
        {
            peers.push(known);
        }
    }
    // Those the recipient would most likely pick as neighbours itself (lowest pair
    // score with it, as `topology` ranks them), if there are too many.
    if peers.len() > MAX_LISTED_PEERS {
        peers.sort_by_cached_key(|known| topology::pair_score(peer, known));
        peers.truncate(MAX_LISTED_PEERS);
    }

    let auth = state
        .swarm_key
        .as_ref()
        .map(|key| key.hello_reply_tag(&swarm_id, state.local_id(), peer))
        .unwrap_or_default();
    let ready = state.is_ready().await;
    let (epoch, epoch_seq) = state.epoch_stamp().await;
    let reply = transport::HelloReply {
        swarm_id,
        peers,
        epoch,
        epoch_seq,
        auth,
        ready,
        run_id,
        away_ms: away.map(|away| away.as_millis() as u64),
        sender_run_id: state.run_id,
    };

    debug!(to = %peer, ready, "Sending HELLO_REPLY");

    state
        .transport
        .send_message(peer, &ParsedMessage::HelloReply(reply))
        .await
        .context("Sending HELLO_REPLY")?;

    trace!(to = %peer, "HELLO_REPLY sent");
    Ok(())
}

/// Greets the peers a HELLO_REPLY listed: at most `MAX_LISTED_PEERS`, only valid
/// identities, and, once the member list is known, only those that are our neighbours
/// (it says whom we greet; a listed peer that isn't one would get nothing from us, so
/// it isn't remembered either).
pub(super) async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[String]) {
    let view_known = topology::view_known(state);
    let mut taken = 0;
    for discovered_peer in advertised_peers {
        if discovered_peer == from_peer || !state.transport.is_peer_id(discovered_peer) {
            continue;
        }
        if view_known && !topology::is_neighbour(state, discovered_peer) {
            continue;
        }
        if taken == MAX_LISTED_PEERS {
            debug!(
                from = %from_peer,
                listed = advertised_peers.len(),
                "HELLO_REPLY listed more peers than we take; ignoring the rest"
            );
            break;
        }
        taken += 1;

        remember_peer(state, discovered_peer).await;
        if let Err(error) = send_hello_once(state, discovered_peer).await {
            warn!(
                peer = %discovered_peer,
                error = %format!("{error:#}"),
                "Failed to initiate HELLO with discovered peer"
            );
        }
    }
}

/// New neighbours get our HELLO (we may not have known them before: the member list is
/// how a partial mesh discovers them) and a sync opening. The opening matters for a
/// neighbour whose handshake is already complete, e.g. one that returned from dormancy
/// while it wasn't our neighbour: it was only answered then, and with no HELLO to send,
/// nothing else would open sync until the keepalive heal (`resync_after`). Former
/// neighbours get no more HELLOs, keepalives or sync openings from us; their session is
/// reset, so we don't keep pushing changes to them either. They are still answered if
/// they contact us.
pub(crate) async fn apply_neighbour_change(state: &AppState, change: topology::NeighbourChange) {
    for peer in &change.added {
        remember_peer(state, peer).await;
        if let Err(error) = send_hello_once(state, peer).await {
            warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to new neighbour failed; will retry");
        }
        // Skipped there unless the handshake is complete; otherwise its reply restarts
        // sync anyway.
        let _ = state
            .sync_events
            .send(SyncEvent::PeerNeedsSync(peer.clone()));
    }
    for peer in &change.removed {
        if let Some(peer_state) = state.get_peer_state(peer).await {
            let mut peer_state = peer_state.lock().await;
            peer_state.stop_hello_retries();
            peer_state.reset_sync();
        }
    }
}

// --- peer membership ---------------------------------------------------------

/// Makes `peer` a member of `swarm_id`, after its valid HELLO or HELLO_REPLY. Either
/// completes the handshake: one round trip proves both sides to each other (the HELLO
/// the sender's key, the reply the swarm ID and the responder's key), so neither side
/// needs a HELLO of its own. `from_hello`: a HELLO that starts a new handshake (a resent
/// or late one doesn't). `ready` is what the message said.
pub(super) async fn establish_peer(
    state: &AppState,
    peer: &str,
    swarm_id: String,
    from_hello: bool,
    ready: bool,
) -> Result<()> {
    let expected_swarm = state.swarm_id().await;

    if let Some(expected) = &expected_swarm {
        if expected != &swarm_id {
            bail!(
                "Incompatible swarm at {peer}: expected {}, received {}",
                expected,
                swarm_id
            );
        }
    }

    let peer_state_mutex = state.get_or_create_peer_state(peer).await;
    let mut peer_state = peer_state_mutex.lock().await;

    if let Standing::Member { swarm_id: existing } = &peer_state.standing {
        if existing != &swarm_id {
            bail!(
                "Peer {peer} changed swarm identity from {} to {}",
                existing,
                swarm_id
            );
        }
    }

    // Lock order: peer mutex -> lifecycle.
    if expected_swarm.is_none() {
        state.set_expected_swarm_id(swarm_id.clone()).await?;
    }

    if !matches!(peer_state.standing, Standing::Member { .. }) {
        peer_state.standing = Standing::Member {
            swarm_id: swarm_id.clone(),
        };
    }

    // Stop our retries and sync from a clean slate. A new handshake from the peer always
    // resets: it starts over with us. Its resent HELLOs and duplicate replies (caused by
    // our own retries) don't.
    peer_state.reports_ready = Some(ready);
    if from_hello || peer_state.hello_state != HelloState::Established {
        peer_state.hello_state = HelloState::Established;
        peer_state.reset_sync();
        info!(peer = %peer, from_hello, "Handshake complete; sync state reset");
    }

    info!(peer = %peer, swarm_id = %swarm_id, "Peer membership established");
    Ok(())
}

#[instrument(skip(state), fields(peer = %peer))]
pub(crate) async fn peer_swarm(state: &AppState, peer: &str) -> Option<String> {
    let peer_state_mutex = state.get_peer_state(peer).await?;
    let peer_state = peer_state_mutex.lock().await;
    peer_state.standing.swarm_id().map(str::to_owned)
}