//! Dispatch of inbound messages.

use std::time::Duration;

use anyhow::{Result, anyhow, bail};
use tokio::time::Instant;
use tracing::{debug, info, warn};

use super::admit::admit;
use super::epoch::observe_peer_epoch;
use super::handshake::{
    discover_peers, ensure_handshake, establish_peer, peer_swarm, remember_peer, send_hello_reply,
};
use super::keepalive::{heal_if_diverged, local_keepalive, note_keepalive};
use super::lanes::SyncLanes;
use crate::state::AppState;
use crate::swarm::sync::SyncEvent;
use crate::transport::{self, ParsedMessage};

/// Handles `message` from `peer` in arrival order, except that a sync message is only
/// queued here, on `lanes`: everything before it (admission, liveness, the peer's epoch)
/// is done in order, applying it isn't.
pub(crate) async fn handle_message(
    state: &AppState,
    lanes: &SyncLanes,
    peer: String,
    message: transport::ParsedMessage,
) -> Result<()> {
    if peer == state.local_id() {
        debug!("Ignoring message from self");
        return Ok(());
    }

    // A reply to a HELLO from an earlier run of ours (late, or queued in Mycelium while
    // we were down) says nothing about now; the peer answers our current HELLO.
    if let transport::ParsedMessage::HelloReply(reply) = &message {
        if reply.run_id != state.run_id {
            state.stats.stale_hello_reply();
            debug!(peer = %peer, "Ignoring a HELLO_REPLY to a HELLO from an earlier run");
            return Ok(());
        }
    }

    if !admit(state, &peer, &message).await {
        return Ok(());
    }

    remember_peer(state, &peer).await;

    // Any inbound message proves liveness. When we heard from it before tells a
    // restarted peer how long it was away (in our HELLO_REPLY).
    let (reactivated, previously_seen) = {
        let peer_state_mutex = state.get_or_create_peer_state(&peer).await;
        let mut peer_state = peer_state_mutex.lock().await;
        let previously_seen = peer_state.last_seen;
        if peer_state.unreachable_since.is_some() {
            // Reachable again (`mark_seen` clears it).
            state.topology_changed.notify_one();
        }
        (peer_state.mark_seen(Instant::now()), previously_seen)
    };
    crate::swarm::source::heard_from(state, &peer);

    if reactivated {
        info!(peer = %peer, "Dormant peer is active again; re-handshaking");
    }

    let run = match &message {
        transport::ParsedMessage::Hello(hello) => Some(hello.run_id),
        transport::ParsedMessage::HelloReply(reply) => Some(reply.sender_run_id),
        _ => None,
    };
    let same_epoch =
        observe_peer_epoch(state, &peer, message.epoch(), message.epoch_seq(), run).await;

    //debug!(from = %peer, message = ?message, "Received message");

    match message {
        transport::ParsedMessage::Hello(hello) => {
            // A HELLO means the peer is (re)starting its handshake with us, e.g. after a
            // restart. It becomes a member right away; our reply makes us one for it.
            let Some(swarm_id) = state.swarm_id().await else {
                // We are joining ourselves; the peer retries until we know our swarm.
                debug!(peer = %peer, "HELLO before we know our swarm; not answering yet");
                return Ok(());
            };
            if let Some(theirs) = &hello.swarm_id {
                if theirs != &swarm_id {
                    bail!("Incompatible swarm at {peer}: ours is {swarm_id}, it claims {theirs}");
                }
            }
            // A resent HELLO (our reply was lost) or a late copy of one repeats a
            // handshake we answered already: it must not reset the session again.
            let new_handshake = {
                let peer_state = state.get_or_create_peer_state(&peer).await;
                let mut peer_state = peer_state.lock().await;
                let new = peer_state.last_hello.is_none_or(|(run, handshake)| {
                    run != hello.run_id || hello.handshake > handshake
                });
                if new {
                    peer_state.last_hello = Some((hello.run_id, hello.handshake));
                }
                new
            };
            establish_peer(state, &peer, swarm_id, new_handshake, hello.ready).await?;

            let away = previously_seen.map(|seen| seen.elapsed());
            if let Err(error) = send_hello_reply(state, &peer, hello.run_id, away).await {
                // The peer retries its HELLO.
                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed");
            }

            let _ = state
                .sync_events
                .send(SyncEvent::PeerValidated(peer.clone()));
        }
        transport::ParsedMessage::HelloReply(reply) => {
            state.lock_rtt().hello_reply_received(&peer, Instant::now());
            crate::swarm::stability::note_away(state, reply.away_ms.map(Duration::from_millis));
            establish_peer(state, &peer, reply.swarm_id.clone(), false, reply.ready).await?;

            discover_peers(state, &peer, &reply.peers).await;

            state
                .sync_events
                .send(SyncEvent::PeerValidated(peer.clone()))
                .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
        }
        transport::ParsedMessage::Sync(envelope) => {
            if reactivated {
                ensure_handshake(state, &peer).await;
            }
            if !same_epoch {
                debug!(peer = %peer, epoch = %envelope.epoch, "Dropping sync message from another epoch");
                return Ok(());
            }
            crate::swarm::source::synced_from(state, &peer);
            lanes.enqueue(&peer, envelope);
        }
        transport::ParsedMessage::Blob(blob) => {
            crate::swarm::blob::handle(state, &peer, blob).await;
        }
        transport::ParsedMessage::Ping(remote) => {
            let local = on_keepalive(state, &peer, &remote, reactivated).await;
            if let Err(error) = state
                .transport
                .send_message(&peer, &ParsedMessage::Pong(local))
                .await
            {
                debug!(peer = %peer, error = %format!("{error:#}"), "PONG failed");
            }
        }
        transport::ParsedMessage::Pong(remote) => {
            // It answers our PING: a round-trip sample.
            state.lock_rtt().pong_received(&peer, Instant::now());
            on_keepalive(state, &peer, &remote, reactivated).await;
        }
    }

    Ok(())
}

/// What a PING and a PONG have in common: a handshake if needed, the peer's keepalive
/// noted, and the heal if our documents differ. Returns our keepalive for the peer.
async fn on_keepalive(
    state: &AppState,
    peer: &str,
    remote: &transport::Keepalive,
    reactivated: bool,
) -> transport::Keepalive {
    if reactivated || peer_swarm(state, peer).await.is_none() {
        ensure_handshake(state, peer).await;
    }
    note_keepalive(state, peer, remote).await;
    let local = local_keepalive(state, peer).await;
    heal_if_diverged(state, peer, &local, remote).await;
    local
}