// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! This module defines the network interface, and provides network implementations for the
//! consensus protocol.
//!
//! Having an abstract network interface allows
//! - simplifying the semantics of sending data and serving requests over the network
//! - hiding implementation specific types and semantics from the consensus protocol
//! - allowing easy swapping of network implementations, for better performance or testing
//!
//! When modifying the client and server interfaces, the principle is to keep the interfaces
//! low level, close to underlying implementations in semantics. For example, the client interface
//! exposes sending messages to a specific peer, instead of broadcasting to all peers. Subscribing
//! to a stream of blocks gets back the stream via response, instead of delivering the stream
//! directly to the server. This keeps the logic agnostics to the underlying network outside of
//! this module, so they can be reused easily across network implementations.

use std::{
    collections::BTreeSet,
    fmt::{Display, Formatter},
    net::SocketAddrV6,
    pin::Pin,
    sync::Arc,
    time::Duration,
};

use async_trait::async_trait;
use bytes::Bytes;
use consensus_config::{AuthorityIndex, NetworkKeyPair, NetworkPublicKey};
use consensus_types::block::{BlockRef, Round};
use fastcrypto::encoding::{Encoding, Hex};
use futures::Stream;
use mysten_network::{Multiaddr, multiaddr::Protocol};
use prost::Message as _;

use crate::{
    block::{ExtendedBlock, VerifiedBlock},
    commit::{CommitRange, TrustedCommit},
    context::Context,
    error::{ConsensusError, ConsensusResult},
};

/// Identifies an observer node by its network public key.
pub type NodeId = NetworkPublicKey;

/// Identifies a peer in the network, which can be either a validator or an observer.
/// The Observer variant is boxed to keep the enum small, since `NodeId` (32 bytes) is
/// much larger than `AuthorityIndex` (4 bytes).
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum PeerId {
    /// A validator node identified by its authority index.
    Validator(AuthorityIndex),
    /// An observer node identified by its network public key.
    Observer(Box<NodeId>),
}

impl Display for PeerId {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            PeerId::Validator(authority) => write!(f, "[{}]", authority),
            PeerId::Observer(node_id) => {
                let bytes = node_id.to_bytes();
                let s = Hex::encode(bytes.get(0..4).ok_or(std::fmt::Error)?);
                write!(f, "o#{}..", s)
            }
        }
    }
}

impl PeerId {
    /// Returns a human-readable name suitable for logging. For observers, prints
    /// the first 8 hex digits of the public key.
    pub(crate) fn hostname(&self, context: &Context) -> String {
        match self {
            PeerId::Validator(index) => context.committee.authority(*index).hostname.to_string(),
            PeerId::Observer(node_id) => {
                let bytes = node_id.to_bytes();
                format!(
                    "[Observer]{:02x}{:02x}{:02x}{:02x}",
                    bytes[0], bytes[1], bytes[2], bytes[3]
                )
            }
        }
    }

    /// Returns a short label suitable for use in metrics. Does not include
    /// full public keys to avoid high-cardinality metric labels.
    pub(crate) fn labelname(&self, context: &Context) -> String {
        match self {
            PeerId::Validator(index) => context.committee.authority(*index).hostname.to_string(),
            PeerId::Observer(_) => "observer".to_string(),
        }
    }
}

// Tonic generated RPC stubs.
mod tonic_gen {
    include!(concat!(env!("OUT_DIR"), "/consensus.ConsensusService.rs"));
    include!(concat!(env!("OUT_DIR"), "/consensus.ObserverService.rs"));
}

mod clients;
pub(crate) mod metrics;
mod metrics_layer;
#[cfg(all(test, not(msim)))]
mod network_tests;
#[cfg(not(msim))]
pub(crate) mod observer;
#[cfg(msim)]
pub mod observer;
#[cfg(test)]
pub(crate) mod test_network;
#[cfg(not(msim))]
pub(crate) mod tonic_network;
#[cfg(msim)]
pub mod tonic_network;
mod tonic_tls;

/// A stream of serialized filtered blocks returned over the network.
pub(crate) type BlockStream = Pin<Box<dyn Stream<Item = ExtendedSerializedBlock> + Send>>;

/// Validator network client for communicating with validator peers.
///
/// NOTE: the timeout parameters help saving resources at client and potentially server.
/// But it is up to the server implementation if the timeout is honored.
/// - To bound server resources, server should implement own timeout for incoming requests.
#[async_trait]
pub(crate) trait ValidatorNetworkClient: Send + Sync + Sized + 'static {
    /// Subscribes to blocks from a peer after last_received round.
    async fn subscribe_blocks(
        &self,
        peer: AuthorityIndex,
        last_received: Round,
        timeout: Duration,
    ) -> ConsensusResult<BlockStream>;

    // TODO: add a parameter for maximum total size of blocks returned.
    /// Fetches serialized `SignedBlock`s from a peer. It also might return additional ancestor blocks
    /// of the requested blocks according to the provided `fetch_after_rounds`. The `fetch_after_rounds`
    /// length should be equal to the committee size. If `fetch_after_rounds` is empty then it will
    /// be simply ignored.
    async fn fetch_blocks(
        &self,
        peer: AuthorityIndex,
        block_refs: Vec<BlockRef>,
        fetch_after_rounds: Vec<Round>,
        fetch_missing_ancestors: bool,
        timeout: Duration,
    ) -> ConsensusResult<Vec<Bytes>>;

    /// Fetches serialized commits in the commit range from a peer.
    /// Returns a tuple of both the serialized commits, and serialized blocks that contain
    /// votes certifying the last commit.
    async fn fetch_commits(
        &self,
        peer: AuthorityIndex,
        commit_range: CommitRange,
        timeout: Duration,
    ) -> ConsensusResult<(Vec<Bytes>, Vec<Bytes>)>;

    /// Fetches the latest block from `peer` for the requested `authorities`. The latest blocks
    /// are returned in the serialised format of `SignedBlocks`. The method can return multiple
    /// blocks per peer as its possible to have equivocations.
    async fn fetch_latest_blocks(
        &self,
        peer: AuthorityIndex,
        authorities: Vec<AuthorityIndex>,
        timeout: Duration,
    ) -> ConsensusResult<Vec<Bytes>>;

    /// Gets the latest received & accepted rounds of all authorities from the peer.
    async fn get_latest_rounds(
        &self,
        peer: AuthorityIndex,
        timeout: Duration,
    ) -> ConsensusResult<(Vec<Round>, Vec<Round>)>;

    /// Sends a serialized SignedBlock to a peer.
    #[cfg(all(test, not(msim)))]
    async fn send_block(
        &self,
        peer: AuthorityIndex,
        block: &VerifiedBlock,
        timeout: Duration,
    ) -> ConsensusResult<()>;
}

/// Validator network service for handling requests from validator peers.
#[async_trait]
pub(crate) trait ValidatorNetworkService: Send + Sync + 'static {
    /// Handles the block sent from the peer via either unicast RPC or subscription stream.
    /// Peer value can be trusted to be a valid authority index.
    /// But serialized_block must be verified before its contents are trusted.
    /// Excluded ancestors are also included as part of an effort to further propagate
    /// blocks to peers despite the current exclusion.
    async fn handle_send_block(
        &self,
        peer: AuthorityIndex,
        block: ExtendedSerializedBlock,
    ) -> ConsensusResult<()>;

    /// Handles the subscription request from the peer.
    /// A stream of newly proposed blocks is returned to the peer.
    /// The stream continues until the end of epoch, peer unsubscribes, or a network error / crash
    /// occurs.
    async fn handle_subscribe_blocks(
        &self,
        peer: AuthorityIndex,
        last_received: Round,
    ) -> ConsensusResult<BlockStream>;

    /// Handles the request to fetch blocks by references from the peer.
    async fn handle_fetch_blocks(
        &self,
        peer: AuthorityIndex,
        block_refs: Vec<BlockRef>,
        fetch_after_rounds: Vec<Round>,
        fetch_missing_ancestors: bool,
    ) -> ConsensusResult<Vec<Bytes>>;

    /// Handles the request to fetch commits by index range from the peer.
    async fn handle_fetch_commits(
        &self,
        peer: AuthorityIndex,
        commit_range: CommitRange,
    ) -> ConsensusResult<(Vec<TrustedCommit>, Vec<VerifiedBlock>)>;

    /// Handles the request to fetch the latest block for the provided `authorities`.
    async fn handle_fetch_latest_blocks(
        &self,
        peer: AuthorityIndex,
        authorities: Vec<AuthorityIndex>,
    ) -> ConsensusResult<Vec<Bytes>>;

    /// Handles the request to get the latest received & accepted rounds of all authorities.
    async fn handle_get_latest_rounds(
        &self,
        peer: AuthorityIndex,
    ) -> ConsensusResult<(Vec<Round>, Vec<Round>)>;
}

/// Handler for randomness round signatures exchanged between validators and
/// observer nodes via the consensus block stream.
pub trait RandomnessSignatureHandler: Send + Sync + 'static {
    /// Called by the observer subscriber for each randomness round signature
    /// received from the block stream.
    fn handle_randomness_signature(&self, data: Bytes);

    /// Returns a receiver for broadcast randomness signatures. Called by
    /// the observer service to merge signatures into the outgoing block stream.
    fn subscribe_randomness_signatures(&self) -> tokio::sync::broadcast::Receiver<Bytes>;
}

/// Filter sent by an observer when opening a block stream, restricting which authorities'
/// blocks the server releases on that stream. Only validated committee members appear here:
/// the transport layer converts and checks the raw wire representation before the filter
/// reaches the service.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub(crate) enum BlockStreamFilter {
    /// Blocks from all authorities are streamed.
    #[default]
    All,
    /// Only blocks authored by the given committee members are streamed. Must be
    /// non-empty; the wire format cannot express an empty author set.
    Authors(BTreeSet<AuthorityIndex>),
}

/// A single item in the observer block stream, carrying both blocks and auxiliary data.
pub(crate) struct ObserverStreamItem {
    pub(crate) blocks: Vec<Bytes>,
    pub(crate) auxiliary_data: observer::AuxiliaryData,
}

/// Observer block stream type.
pub(crate) type ObserverBlockStream =
    Pin<Box<dyn Stream<Item = ObserverStreamItem> + Send + 'static>>;

/// Observer network service for handling requests from observer nodes.
/// Unlike ValidatorNetworkService which uses AuthorityIndex, this uses NodeId (NetworkPublicKey)
/// to identify peers since observers are not part of the committee.
#[async_trait]
pub(crate) trait ObserverNetworkService: Send + Sync + 'static {
    /// Handles a block received from a peer subscription. Used by ObserverSubscriber to process
    /// blocks streamed from validators or other observers.
    async fn handle_block(&self, peer: PeerId, block: Bytes) -> ConsensusResult<()>;

    /// Handles the block streaming request from an observer peer.
    /// Returns a stream of blocks with the highest commit index for each block.
    /// Blocks with rounds higher than the highest_round_per_authority will be streamed,
    /// restricted to the authors requested in the filter.
    async fn handle_stream_blocks(
        &self,
        peer: NodeId,
        highest_round_per_authority: Vec<Round>,
        filter: BlockStreamFilter,
    ) -> ConsensusResult<ObserverBlockStream>;

    /// Handles the request to fetch blocks by references from an observer peer.
    async fn handle_fetch_blocks(
        &self,
        peer: NodeId,
        block_refs: Vec<BlockRef>,
        fetch_after_rounds: Vec<Round>,
        fetch_missing_ancestors: bool,
    ) -> ConsensusResult<Vec<Bytes>>;

    /// Handles the request to fetch commits by index range from an observer peer.
    /// Returns serialized commits and certifier blocks.
    async fn handle_fetch_commits(
        &self,
        peer: NodeId,
        commit_range: CommitRange,
    ) -> ConsensusResult<(Vec<TrustedCommit>, Vec<VerifiedBlock>)>;
}

/// Observer network client for communicating with validators' observer ports or other observers.
/// Unlike ValidatorNetworkClient which uses AuthorityIndex, this uses PeerId to identify peers
/// since the observer server can serve both validators and observer nodes.
#[async_trait]
pub(crate) trait ObserverNetworkClient: Send + Sync + Sized + 'static {
    /// Initiates block streaming with a peer (validator or observer).
    /// Returns a stream of blocks with the highest commit index.
    /// Blocks with rounds higher than the highest_round_per_authority will be streamed,
    /// restricted to the authors requested in the filter.
    async fn stream_blocks(
        &self,
        peer: PeerId,
        highest_round_per_authority: Vec<Round>,
        filter: BlockStreamFilter,
        timeout: Duration,
    ) -> ConsensusResult<ObserverBlockStream>;

    /// Fetches serialized blocks by references from a peer.
    async fn fetch_blocks(
        &self,
        peer: PeerId,
        block_refs: Vec<BlockRef>,
        fetch_after_rounds: Vec<Round>,
        fetch_missing_ancestors: bool,
        timeout: Duration,
    ) -> ConsensusResult<Vec<Bytes>>;

    /// Fetches serialized commits in the commit range from a peer.
    /// Returns a tuple of both the serialized commits, and serialized blocks that contain
    /// votes certifying the last commit.
    async fn fetch_commits(
        &self,
        peer: PeerId,
        commit_range: CommitRange,
        timeout: Duration,
    ) -> ConsensusResult<(Vec<Bytes>, Vec<Bytes>)>;
}

/// An `AuthorityNode` holds a `NetworkManager` until shutdown.
/// Dropping `NetworkManager` will shutdown the network service.
pub(crate) trait NetworkManager: Send + Sync {
    type ValidatorClient: ValidatorNetworkClient;
    type ObserverClient: ObserverNetworkClient;

    /// Creates a new network manager.
    fn new(context: Arc<Context>, network_keypair: NetworkKeyPair) -> Self;

    /// Returns the validator network client.
    fn validator_client(&self) -> Arc<Self::ValidatorClient>;

    /// Returns the observer network client.
    fn observer_client(&self) -> Arc<Self::ObserverClient>;

    /// Starts the validator network server with the provided service.
    async fn start_validator_server<V>(&mut self, service: Arc<V>)
    where
        V: ValidatorNetworkService;

    /// Starts the observer network server with the provided service.
    async fn start_observer_server<O>(&mut self, service: Arc<O>)
    where
        O: ObserverNetworkService;

    /// Stops the network service.
    async fn stop(&mut self);

    /// Updates the network address for a peer identified by their authority index.
    /// If address is None, the override is cleared and the committee address will be used.
    fn update_peer_address(&self, peer: AuthorityIndex, address: Option<Multiaddr>);
}

// Re-export the concrete client implementations.
pub(crate) use clients::{CommitSyncerClient, SynchronizerClient};

/// Serialized block with extended information from the proposing authority.
#[derive(Clone, PartialEq, Eq, Debug)]
pub(crate) struct ExtendedSerializedBlock {
    pub(crate) block: SerializedBlockForm,
    // Serialized BlockRefs that are excluded from the blocks ancestors.
    pub(crate) excluded_ancestors: Vec<Vec<u8>>,
}

/// Wire framing for one block on live subscription streams, carried inside a bytes
/// field so batch streams can reuse it.
#[derive(Clone, PartialEq, prost::Message)]
pub(crate) struct SerializedBlockEnvelope {
    #[prost(oneof = "SerializedBlockForm", tags = "1, 2")]
    pub(crate) block: Option<SerializedBlockForm>,
}

/// Which representation of a block a payload holds.
#[derive(Clone, PartialEq, Eq, prost::Oneof)]
pub(crate) enum SerializedBlockForm {
    /// The block exactly as serialized and signed by its author.
    #[prost(bytes = "bytes", tag = "1")]
    Full(Bytes),
    /// The same block with its ancestor digests stripped, to be rebuilt by the
    /// receiver. Nothing emits this form yet; the encoder lands with the codec.
    #[prost(bytes = "bytes", tag = "2")]
    Slim(Bytes),
}

impl SerializedBlockEnvelope {
    pub(crate) fn encode_form(form: SerializedBlockForm) -> Bytes {
        Self { block: Some(form) }.encode_to_vec().into()
    }

    pub(crate) fn decode_form(bytes: &[u8]) -> ConsensusResult<SerializedBlockForm> {
        Self::decode(bytes)
            .map_err(ConsensusError::MalformedBlockEnvelope)?
            .block
            .ok_or_else(|| {
                ConsensusError::MalformedBlockEnvelope(prost::DecodeError::new(
                    "empty block envelope",
                ))
            })
    }
}

impl From<ExtendedBlock> for ExtendedSerializedBlock {
    fn from(extended_block: ExtendedBlock) -> Self {
        Self {
            block: SerializedBlockForm::Full(extended_block.block.serialized().clone()),
            excluded_ancestors: extended_block
                .excluded_ancestors
                .iter()
                .filter_map(|r| match bcs::to_bytes(r) {
                    Ok(serialized) => Some(serialized),
                    Err(e) => {
                        tracing::debug!("Failed to serialize block ref {:?}: {e:?}", r);
                        None
                    }
                })
                .collect(),
        }
    }
}

/// Attempts to convert a multiaddr of the form `/[ip4,ip6,dns]/{}/udp/{port}` into
/// a host:port string.
pub(crate) fn to_host_port_str(addr: &Multiaddr) -> Result<String, String> {
    let mut iter = addr.iter();

    match (iter.next(), iter.next()) {
        (Some(Protocol::Ip4(ipaddr)), Some(Protocol::Udp(port))) => {
            Ok(format!("{}:{}", ipaddr, port))
        }
        (Some(Protocol::Ip6(ipaddr)), Some(Protocol::Udp(port))) => {
            Ok(format!("{}", SocketAddrV6::new(ipaddr, port, 0, 0)))
        }
        (Some(Protocol::Dns(hostname)), Some(Protocol::Udp(port))) => {
            Ok(format!("{}:{}", hostname, port))
        }

        _ => Err(format!("unsupported multiaddr: {addr}")),
    }
}