use anyhow::Result;
use async_trait::async_trait;
use automerge::ChangeHash;
use automerge::sync::Message;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::boxed::Box;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};

pub(crate) use crate::document::replica::Epoch;

pub(crate) mod firewall;
pub(crate) mod mycelium;
#[cfg(test)]
pub(crate) mod sim;
pub(crate) mod traffic;
pub(crate) mod wire;

/// Random identifier of one run of a node (from start to exit).
pub(crate) type RunId = [u8; 16];

/// Most heads listed in a keepalive; more are sent as a digest. A head takes 34 bytes on
/// the wire, so 32 of them keep a keepalive within one Mycelium chunk (1300 bytes; see
/// the size test in `mycelium`). There are rarely more heads than nodes that wrote
/// concurrently since the last merging write.
pub(crate) const MAX_LISTED_HEADS: usize = 32;

/// A document's heads, compactly: the heads themselves while there are few, otherwise a
/// digest of them with some of them. Two summaries are equal exactly when the heads are,
/// whatever form each took (nodes with different limits still agree).
#[derive(Debug, Clone)]
pub(crate) enum HeadsSummary {
    /// Sorted.
    Heads(Vec<ChangeHash>),
    Digest {
        /// SHA-256 over the sorted heads.
        digest: [u8; 32],
        /// `MAX_LISTED_HEADS` of the heads, another window of them each time: a receiver
        /// pulls those it lacks, and over a few keepalives every head is listed, instead
        /// of starting the sync over to find out what it lacks.
        sample: Vec<ChangeHash>,
    },
}

/// Where the next sample starts (`HeadsSummary::of`).
static SAMPLE_OFFSET: AtomicUsize = AtomicUsize::new(0);

impl HeadsSummary {
    pub(crate) fn of(mut heads: Vec<ChangeHash>) -> Self {
        heads.sort_unstable();
        if heads.len() <= MAX_LISTED_HEADS {
            return Self::Heads(heads);
        }
        let start = SAMPLE_OFFSET.fetch_add(MAX_LISTED_HEADS, Ordering::Relaxed) % heads.len();
        let sample = heads
            .iter()
            .cycle()
            .skip(start)
            .take(MAX_LISTED_HEADS)
            .copied()
            .collect();
        Self::Digest {
            digest: digest_of(&heads),
            sample,
        }
    }

    /// The digest of the heads, computed if they are listed.
    pub(crate) fn digest(&self) -> [u8; 32] {
        match self {
            Self::Heads(heads) => digest_of(heads),
            Self::Digest { digest, .. } => *digest,
        }
    }

    /// The heads, if all are listed.
    pub(crate) fn heads(&self) -> Option<&[ChangeHash]> {
        match self {
            Self::Heads(heads) => Some(heads),
            Self::Digest { .. } => None,
        }
    }

    /// The heads listed: all of them, or a sample.
    pub(crate) fn listed(&self) -> &[ChangeHash] {
        match self {
            Self::Heads(heads) => heads,
            Self::Digest { sample, .. } => sample,
        }
    }
}

impl PartialEq for HeadsSummary {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Heads(a), Self::Heads(b)) => a == b,
            _ => self.digest() == other.digest(),
        }
    }
}

impl Eq for HeadsSummary {}

/// SHA-256 over sorted heads.
fn digest_of(sorted_heads: &[ChangeHash]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    for head in sorted_heads {
        hasher.update(head.0);
    }
    hasher.finalize().into()
}

/// PING and PONG. Addressed to one peer: `syncs_with_you` is about the recipient.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct Keepalive {
    /// Only comparable between nodes on the same epoch.
    #[serde(with = "crate::transport::wire::codec::heads")]
    pub(crate) heads: HeadsSummary,
    #[serde(with = "crate::transport::wire::codec::epoch")]
    pub(crate) epoch: Epoch,
    /// See `SyncEnvelope::epoch_seq`.
    pub(crate) epoch_seq: u32,
    /// The sender is ready (same meaning as `SyncEnvelope::ready`).
    pub(crate) ready: bool,
    /// The sender would sync with the recipient: counts it as an active member on its
    /// own epoch.
    pub(crate) syncs_with_you: bool,
}

#[derive(Debug)]
pub(crate) enum ParsedMessage {
    Hello(HelloRequest),
    HelloReply(HelloReply),
    Sync(SyncEnvelope),
    Ping(Keepalive),
    Pong(Keepalive),
    Blob(BlobMessage),
}

/// Content-addressed bytes that any node holding them can serve, in chunks
/// (`swarm::blob`).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub(crate) enum BlobId {
    /// An Automerge change, by its hash: e.g. an epoch's snapshot change.
    Change(#[serde(with = "crate::transport::wire::codec::change_hash")] ChangeHash),
    /// Bytes by their BLAKE3 hash: e.g. the content of a remake.
    Content(#[serde(with = "crate::transport::wire::codec::bytes32")] [u8; 32]),
    /// The changes of a `DeltaSpec`, by its ID (`DeltaSpec::id`). The requests carry
    /// the spec, so any node holding its `to` can serve it.
    Delta(#[serde(with = "crate::transport::wire::codec::bytes32")] [u8; 32]),
}

impl std::fmt::Display for BlobId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BlobId::Change(hash) => write!(f, "change:{hash}"),
            BlobId::Content(hash) => write!(f, "content:{}", Hex(hash)),
            BlobId::Delta(id) => write!(f, "delta:{}", Hex(id)),
        }
    }
}

/// A blob request or one of its answers. Each is a message of its own, so no message
/// is larger than one chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct BlobMessage {
    /// The sender's epoch (every message carries it).
    #[serde(with = "crate::transport::wire::codec::epoch")]
    pub(crate) epoch: Epoch,
    /// See `SyncEnvelope::epoch_seq`.
    pub(crate) epoch_seq: u32,
    pub(crate) id: BlobId,
    pub(crate) body: BlobBody,
}

#[derive(Clone, Serialize, Deserialize)]
pub(crate) enum BlobBody {
    /// Asks for these chunks; for the manifest if empty. `delta`: the spec of a
    /// `BlobId::Delta`, which the ID alone doesn't say.
    Request {
        chunks: Vec<u32>,
        delta: Option<DeltaSpec>,
    },
    /// Size and per-chunk BLAKE3 hashes; the last chunk may be shorter.
    Manifest {
        size: u64,
        chunk_size: u32,
        #[serde(with = "crate::transport::wire::codec::hashes32")]
        hashes: Vec<[u8; 32]>,
    },
    Chunk {
        index: u32,
        #[serde(with = "crate::transport::wire::codec::shared_bytes")]
        data: bytes::Bytes,
    },
    /// The sender doesn't have the blob.
    Missing,
}

/// Sizes instead of contents: logged messages stay short.
impl std::fmt::Debug for BlobBody {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BlobBody::Request { chunks, delta } => f
                .debug_struct("Request")
                .field("chunks", chunks)
                .field("delta", &delta.is_some())
                .finish(),
            BlobBody::Manifest {
                size,
                chunk_size,
                hashes,
            } => f
                .debug_struct("Manifest")
                .field("size", size)
                .field("chunk_size", chunk_size)
                .field("chunks", &hashes.len())
                .finish(),
            BlobBody::Chunk { index, data } => f
                .debug_struct("Chunk")
                .field("index", index)
                .field("bytes", &data.len())
                .finish(),
            BlobBody::Missing => f.write_str("Missing"),
        }
    }
}

/// What a message is, without its contents.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum MessageKind {
    Hello,
    HelloReply,
    Sync,
    Ping,
    Pong,
    Blob,
}

impl ParsedMessage {
    pub(crate) fn kind(&self) -> MessageKind {
        match self {
            ParsedMessage::Hello(_) => MessageKind::Hello,
            ParsedMessage::HelloReply(_) => MessageKind::HelloReply,
            ParsedMessage::Sync(_) => MessageKind::Sync,
            ParsedMessage::Ping(_) => MessageKind::Ping,
            ParsedMessage::Pong(_) => MessageKind::Pong,
            ParsedMessage::Blob(_) => MessageKind::Blob,
        }
    }

    /// Every message carries the sender's epoch.
    pub(crate) fn epoch(&self) -> Epoch {
        match self {
            ParsedMessage::Hello(hello) => hello.epoch,
            ParsedMessage::HelloReply(reply) => reply.epoch,
            ParsedMessage::Sync(envelope) => envelope.epoch,
            ParsedMessage::Ping(keepalive) | ParsedMessage::Pong(keepalive) => keepalive.epoch,
            ParsedMessage::Blob(blob) => blob.epoch,
        }
    }

    /// `SyncEnvelope::epoch_seq`, as every message carries it.
    pub(crate) fn epoch_seq(&self) -> u32 {
        match self {
            ParsedMessage::Hello(hello) => hello.epoch_seq,
            ParsedMessage::HelloReply(reply) => reply.epoch_seq,
            ParsedMessage::Sync(envelope) => envelope.epoch_seq,
            ParsedMessage::Ping(keepalive) | ParsedMessage::Pong(keepalive) => keepalive.epoch_seq,
            ParsedMessage::Blob(blob) => blob.epoch_seq,
        }
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub(crate) struct HelloRequest {
    /// The sender's swarm, if it knows it (a new node doesn't yet). A receiver in
    /// another swarm doesn't answer.
    pub(crate) swarm_id: Option<String>,
    #[serde(with = "crate::transport::wire::codec::epoch")]
    pub(crate) epoch: Epoch,
    /// See `SyncEnvelope::epoch_seq`.
    pub(crate) epoch_seq: u32,
    /// Proof of the swarm key (`swarm::auth`); empty without one.
    #[serde(with = "serde_bytes")]
    pub(crate) auth: Vec<u8>,
    /// The sender is ready (same meaning as `SyncEnvelope::ready`).
    pub(crate) ready: bool,
    /// Random per run of the sender; the reply echoes it.
    #[serde(with = "crate::transport::wire::codec::run_id")]
    pub(crate) run_id: RunId,
    /// The handshake's number: the sender numbers the handshakes it starts in a run,
    /// with all peers, from 1. A retry carries the same number, so the recipient resets
    /// its sync session only for a new handshake, not for every resend or late copy.
    pub(crate) handshake: u32,
}

#[derive(Clone, Serialize, Deserialize)]
pub(crate) struct HelloReply {
    pub(crate) swarm_id: String,
    pub(crate) peers: Vec<String>,
    #[serde(with = "crate::transport::wire::codec::epoch")]
    pub(crate) epoch: Epoch,
    /// See `SyncEnvelope::epoch_seq`.
    pub(crate) epoch_seq: u32,
    /// Proof of the swarm key (`swarm::auth`); empty without one.
    #[serde(with = "serde_bytes")]
    pub(crate) auth: Vec<u8>,
    /// The sender is ready.
    pub(crate) ready: bool,
    /// The `run_id` of the HELLO answered. A reply to a HELLO from an earlier run of the
    /// recipient (delivered late, or queued while it was down) is ignored.
    #[serde(with = "crate::transport::wire::codec::run_id")]
    pub(crate) run_id: RunId,
    /// How long before this HELLO the sender had last heard from the recipient, if ever
    /// (in milliseconds): tells a restarted node how long it was away.
    pub(crate) away_ms: Option<u64>,
    /// The sender's own run (as its HELLOs carry it): a new one means it restarted.
    #[serde(with = "crate::transport::wire::codec::run_id")]
    pub(crate) sender_run_id: RunId,
}

/// Bytes shown as lowercase hex, e.g. run IDs in logs.
pub(crate) struct Hex<'a>(pub(crate) &'a [u8]);

impl std::fmt::Display for Hex<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.iter().try_for_each(|byte| write!(f, "{byte:02x}"))
    }
}

impl std::fmt::Debug for Hex<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(self, f)
    }
}

/// Whether a handshake carries a swarm-key proof, not the proof: logged messages don't
/// hand out tags, and stay short.
fn auth_shown(auth: &[u8]) -> &'static str {
    if auth.is_empty() { "none" } else { "present" }
}

impl std::fmt::Debug for HelloRequest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HelloRequest")
            .field("swarm_id", &self.swarm_id)
            .field("epoch", &self.epoch)
            .field("epoch_seq", &self.epoch_seq)
            .field("auth", &auth_shown(&self.auth))
            .field("ready", &self.ready)
            .field("run_id", &Hex(&self.run_id))
            .field("handshake", &self.handshake)
            .finish()
    }
}

impl std::fmt::Debug for HelloReply {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HelloReply")
            .field("swarm_id", &self.swarm_id)
            .field("peers", &self.peers.len())
            .field("epoch", &self.epoch)
            .field("epoch_seq", &self.epoch_seq)
            .field("auth", &auth_shown(&self.auth))
            .field("ready", &self.ready)
            .field("run_id", &Hex(&self.run_id))
            .field("away_ms", &self.away_ms)
            .field("sender_run_id", &Hex(&self.sender_run_id))
            .finish()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct SyncEnvelope {
    pub(crate) swarm_id: String,
    /// Epoch of the document the message was generated from.
    #[serde(with = "crate::transport::wire::codec::epoch")]
    pub(crate) epoch: Epoch,
    /// How many times the sender changed epoch in its current run, when it sent this
    /// (`Replica::epoch_seq`). Messages may arrive late and out of order: the receiver
    /// ignores an epoch older than the newest it has seen in the same session.
    pub(crate) epoch_seq: u32,
    pub(crate) body: SyncBody,
    /// The sender was ready when it generated the message: its document is the swarm's,
    /// not one it is still filling (joining, rejoining). Only such a peer can tell a
    /// catching-up node that it has everything; an empty document says nothing else.
    pub(crate) ready: bool,
}

/// What a sync envelope carries.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) enum SyncBody {
    /// An Automerge sync message.
    Message(#[serde(with = "crate::transport::wire::codec::sync_message")] SyncPayload),
    /// Instead of a message too large to send in one piece: fetch these changes as a
    /// blob (`BlobId::Delta`), then sync on.
    Delta(DeltaSpec),
}

/// The changes in the history of `to` but not of `from`: what a peer at `from` lacks of
/// a document at `to`. Every node holding `to` produces the same bytes for it
/// (`swarm::blob::delta_bytes`).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct DeltaSpec {
    /// Sorted.
    #[serde(with = "crate::transport::wire::codec::change_hashes")]
    pub(crate) from: Vec<ChangeHash>,
    /// Sorted.
    #[serde(with = "crate::transport::wire::codec::change_hashes")]
    pub(crate) to: Vec<ChangeHash>,
}

impl DeltaSpec {
    pub(crate) fn new(mut from: Vec<ChangeHash>, mut to: Vec<ChangeHash>) -> Self {
        from.sort_unstable();
        from.dedup();
        to.sort_unstable();
        to.dedup();
        Self { from, to }
    }

    /// Names the spec as a blob: SHA-256 over both lists.
    pub(crate) fn id(&self) -> [u8; 32] {
        let mut hasher = Sha256::new();
        hasher.update(b"continuum/delta/1");
        for heads in [&self.from, &self.to] {
            hasher.update((heads.len() as u32).to_be_bytes());
            for head in heads {
                hasher.update(head.0);
            }
        }
        hasher.finalize().into()
    }
}

/// An Automerge sync message in Automerge's own encoding. Decoded only where it is
/// applied (`decode`), after the sender proved to be a member and off the consumer: a
/// large message from anyone else costs nothing but its bytes.
#[derive(Clone)]
pub(crate) struct SyncPayload(pub(crate) bytes::Bytes);

impl SyncPayload {
    pub(crate) fn encode(message: Message) -> Self {
        Self(message.encode().into())
    }

    pub(crate) fn decode(&self) -> Result<Message> {
        Message::decode(&self.0)
            .map_err(|error| anyhow::anyhow!("Automerge decode failed: {error:?}"))
    }
}

/// Sizes instead of contents: logged messages stay short.
impl std::fmt::Debug for SyncPayload {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SyncPayload({} bytes)", self.0.len())
    }
}

/// A message serialized for the wire. Cloning is cheap (shared buffer).
#[derive(Clone, Debug)]
pub(crate) struct EncodedMessage {
    pub(crate) peer: String,
    pub(crate) body: bytes::Bytes,
    pub(crate) kind: MessageKind,
    /// Bytes the message takes on the wire (without the transport's own framing).
    pub(crate) payload_len: usize,
}

/// Called with the sender, the message and its `payload_len` (as in `EncodedMessage`).
pub(crate) type HandlerFn = dyn Fn(String, ParsedMessage, usize) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>
    + Send
    + Sync
    + 'static;
pub(crate) type MessageHandler = Box<HandlerFn>;

#[async_trait]
pub(crate) trait Transport: Send + Sync {
    /// This node's identity on the transport (for Mycelium: its public key).
    fn local_id(&self) -> &str;

    /// Whether `id` can be a node's identity on this transport (for Mycelium: a public
    /// key). Identities other peers name to us (in a HELLO_REPLY) are checked with it.
    fn is_peer_id(&self, id: &str) -> bool;

    /// How long the transport tries to deliver a message before dropping it: after that,
    /// a message sent is delivered or lost for good.
    fn delivery_limit(&self) -> std::time::Duration;

    /// Resolves a seed (anything that leads to a node, e.g. an address) to that node's
    /// identity, as `local_id` gives it on that node.
    async fn resolve_seed(&self, seed: &str) -> Result<String>;

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage>;
    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()>;

    /// Sends an outbound message payload to a targeted peer on a specific topic.
    async fn send_message(&self, peer: &str, message: &ParsedMessage) -> Result<()> {
        let encoded = self.encode(peer, message)?;
        self.send_encoded(&encoded).await
    }

    /// Starts the long-polling inbound event loops, driving messages to the provided handler closure.
    async fn start_consumer(&self, handler: MessageHandler) -> Result<()>;

    /// Asynchronously monitors for fatal network client or sub-process failures.
    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::testrand::for_seeds;

    /// Two summaries are equal exactly when the heads are, in any order, whichever form
    /// (listed or digest) each took.
    #[test]
    fn heads_summary_equality_is_set_equality() {
        for_seeds(300, |rng| {
            let n = rng.between(0, 2 * MAX_LISTED_HEADS + 2);
            let heads: Vec<ChangeHash> = (0..n).map(|_| ChangeHash(rng.bytes32())).collect();
            let summary = HeadsSummary::of(heads.clone());
            assert_eq!(summary.heads().is_some(), n <= MAX_LISTED_HEADS);

            let mut shuffled = heads.clone();
            rng.shuffle(&mut shuffled);
            assert_eq!(summary, HeadsSummary::of(shuffled));
            assert_eq!(
                summary,
                HeadsSummary::Digest {
                    digest: summary.digest(),
                    sample: Vec::new(),
                }
            );
            // A sample is some of the heads.
            assert!(summary.listed().iter().all(|head| heads.contains(head)));
            assert_eq!(summary.listed().len(), n.min(MAX_LISTED_HEADS));

            let mut other = heads.clone();
            match rng.below(3) {
                1 if !other.is_empty() => {
                    let i = rng.below(other.len());
                    other.remove(i);
                }
                2 if !other.is_empty() => {
                    let i = rng.below(other.len());
                    other[i] = ChangeHash(rng.bytes32());
                }
                _ => other.push(ChangeHash(rng.bytes32())),
            }
            assert_ne!(summary, HeadsSummary::of(other));
        });
    }

    /// Successive samples cover every head.
    #[test]
    fn samples_rotate_over_all_heads() {
        let heads: Vec<ChangeHash> = (0..100u8).map(|n| ChangeHash([n; 32])).collect();
        let mut seen = std::collections::HashSet::new();
        for _ in 0..heads.len().div_ceil(MAX_LISTED_HEADS) + 1 {
            seen.extend(HeadsSummary::of(heads.clone()).listed().iter().copied());
        }
        // Other tests sample concurrently and move the window too, so allow a few more.
        for _ in 0..100 {
            if seen.len() == heads.len() {
                break;
            }
            seen.extend(HeadsSummary::of(heads.clone()).listed().iter().copied());
        }
        assert_eq!(seen.len(), heads.len());
    }
}