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

use std::{
    collections::BTreeMap,
    net::{SocketAddr, SocketAddrV4, SocketAddrV6},
    pin::Pin,
    sync::{Arc, Weak},
    time::{Duration, Instant},
};

use async_trait::async_trait;
use bytes::Bytes;
use consensus_config::{AuthorityIndex, NetworkKeyPair, NetworkPublicKey};
use consensus_types::block::{BlockRef, Round};
use fastcrypto::{encoding::Encoding, traits::ToFromBytes};
use futures::{Stream, StreamExt as _, stream};
use mysten_network::{Multiaddr, multiaddr::Protocol};
use parking_lot::RwLock;
use sui_http::{
    ServerHandle,
    middleware::{
        callback::{CallbackLayer, MakeCallbackHandler, RequestBody, ResponseHandler},
        grpc_timeout::GrpcTimeout,
    },
};
use sui_tls::AllowPublicKeys;
use tokio_stream::{Iter, iter};
use tonic::{Request, Response, Streaming, codec::CompressionEncoding};
use tower_http::trace::{DefaultMakeSpan, DefaultOnFailure, TraceLayer};
use tracing::{debug, error, info, trace, warn};

use super::{
    BlockStream, ExtendedSerializedBlock, NetworkManager, ObserverNetworkService,
    SerializedBlockEnvelope, SerializedBlockForm, ValidatorNetworkClient, ValidatorNetworkService,
    metrics_layer::{MetricsCallbackMaker, MetricsResponseCallback, SizedRequest, SizedResponse},
    observer::{ObserverPeerInfo, ObserverServiceProxy, TonicObserverClient},
    tonic_gen::{
        consensus_service_client::ConsensusServiceClient,
        consensus_service_server::ConsensusService,
    },
};
use crate::{
    CommitIndex,
    commit::CommitRange,
    context::Context,
    error::{ConsensusError, ConsensusResult},
    network::{
        to_host_port_str,
        tonic_gen::{
            consensus_service_server::ConsensusServiceServer,
            observer_service_server::ObserverServiceServer,
        },
        tonic_tls::certificate_server_name,
    },
};

// Maximum bytes size in a single fetch_blocks()response.
// TODO: put max RPC response size in protocol config.
pub(crate) const MAX_FETCH_RESPONSE_BYTES: usize = 4 * 1024 * 1024;

pub(crate) fn max_fetch_blocks_response_bytes(
    context: &Context,
    block_refs: &[BlockRef],
    fetch_after_rounds: &[Round],
) -> usize {
    // Mirror the server-side limit selection in `block_sync_service::fetch_blocks`:
    // live sync (both inputs non-empty) is capped by `max_blocks_per_sync`,
    // every other mode by `max_blocks_per_fetch`.
    let max_response_num_blocks = if !fetch_after_rounds.is_empty() && !block_refs.is_empty() {
        context.parameters.max_blocks_per_sync
    } else {
        context.parameters.max_blocks_per_fetch
    };

    max_response_num_blocks
        .saturating_mul(context.protocol_config.max_transactions_in_block_bytes() as usize)
        .saturating_mul(2)
}

const DEFAULT_GRPC_SERVER_TIMEOUT: Duration = Duration::from_secs(300);

// HTTP/2 connection and stream window sizes for both validator and observer servers.
const HTTP2_INITIAL_CONNECTION_WINDOW_SIZE: u32 = 64 << 20; // 64 MB
const HTTP2_INITIAL_STREAM_WINDOW_SIZE: u32 = 32 << 20; // 32 MB

// Implements Tonic RPC client for validator consensus operations.
pub(crate) struct TonicValidatorClient {
    context: Arc<Context>,
    network_keypair: NetworkKeyPair,
    channel_pool: Arc<ChannelPool>,
}

impl TonicValidatorClient {
    pub(crate) fn new(context: Arc<Context>, network_keypair: NetworkKeyPair) -> Self {
        Self {
            context: context.clone(),
            network_keypair,
            channel_pool: Arc::new(ChannelPool::new(context)),
        }
    }

    async fn get_client(
        &self,
        peer: AuthorityIndex,
        timeout: Duration,
    ) -> ConsensusResult<ConsensusServiceClient<Channel>> {
        let config = &self.context.parameters.tonic;
        let channel = self
            .channel_pool
            .get_channel(self.network_keypair.clone(), peer, timeout)
            .await?;
        let client = ConsensusServiceClient::new(channel)
            .max_encoding_message_size(config.message_size_limit)
            .max_decoding_message_size(config.message_size_limit)
            .send_compressed(CompressionEncoding::Zstd)
            .accept_compressed(CompressionEncoding::Zstd);
        Ok(client)
    }

    pub(crate) fn update_peer_address(&self, peer: AuthorityIndex, address: Option<Multiaddr>) {
        self.channel_pool.update_address(peer, address);
    }
}

// TODO: make sure callsites do not send request to own index, and return error otherwise.
#[async_trait]
impl ValidatorNetworkClient for TonicValidatorClient {
    async fn subscribe_blocks(
        &self,
        peer: AuthorityIndex,
        last_received: Round,
        timeout: Duration,
    ) -> ConsensusResult<BlockStream> {
        let mut client = self.get_client(peer, timeout).await?;
        // TODO: add sampled block acknowledgments for latency measurements.
        let request = Request::new(stream::once(async move {
            SubscribeBlocksRequest {
                last_received_round: last_received,
            }
        }));
        let response = client.subscribe_blocks(request).await.map_err(|e| {
            ConsensusError::NetworkRequest(format!("subscribe_blocks failed: {e:?}"))
        })?;
        let slim_enabled = self
            .context
            .protocol_config
            .slim_block_propagation_enabled();
        let metrics_context = self.context.clone();
        let peer_hostname = self.context.committee.authority(peer).hostname.clone();
        let stream = response
            .into_inner()
            .map(move |b| match b {
                Ok(response) => {
                    match interpret_subscription_block(response.block, slim_enabled) {
                        Ok(block) => Some(ExtendedSerializedBlock {
                            block,
                            excluded_ancestors: response.excluded_ancestors,
                        }),
                        // A payload that does not frame as an envelope is a peer or
                        // configuration fault, not an invalid block. Ending the stream
                        // here, before anything is delivered, keeps the subscriber's
                        // retry counter and backoff unreset, so a persistently
                        // mismatched peer escalates through reconnect backoff instead
                        // of being rejected block by block forever.
                        Err(e) => {
                            debug!("Failed to decode block envelope from {}: {e:?}", peer);
                            let reason: &'static str = (&e).into();
                            metrics_context
                                .metrics
                                .node_metrics
                                .subscribe_stream_form_failures
                                .with_label_values(&[peer_hostname.as_str(), reason])
                                .inc();
                            None
                        }
                    }
                }
                Err(e) => {
                    debug!("Network error received from {}: {e:?}", peer);
                    None
                }
            })
            .take_while(|item| futures::future::ready(item.is_some()))
            .map(|item| item.expect("terminated by take_while"));
        let rate_limited_stream =
            tokio_stream::StreamExt::throttle(stream, self.context.parameters.min_round_delay / 2)
                .boxed();
        Ok(rate_limited_stream)
    }

    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>> {
        let mut client = self.get_client(peer, timeout).await?;
        let max_allowed_bytes =
            max_fetch_blocks_response_bytes(&self.context, &block_refs, &fetch_after_rounds);
        let mut request = Request::new(FetchBlocksRequest {
            block_refs: block_refs
                .iter()
                .filter_map(|r| match bcs::to_bytes(r) {
                    Ok(serialized) => Some(serialized),
                    Err(e) => {
                        debug!("Failed to serialize block ref {:?}: {e:?}", r);
                        None
                    }
                })
                .collect(),
            fetch_after_rounds,
            fetch_missing_ancestors,
        });
        request.set_timeout(timeout);

        let mut stream = client
            .fetch_blocks(request)
            .await
            .map_err(|e| {
                if e.code() == tonic::Code::DeadlineExceeded {
                    ConsensusError::NetworkRequestTimeout(format!("fetch_blocks failed: {e:?}"))
                } else {
                    ConsensusError::NetworkRequest(format!("fetch_blocks failed: {e:?}"))
                }
            })?
            .into_inner();

        let mut blocks = vec![];
        let mut total_fetched_bytes = 0;
        loop {
            match stream.message().await {
                Ok(Some(response)) => {
                    for b in &response.blocks {
                        total_fetched_bytes += b.len();
                    }
                    blocks.extend(response.blocks);
                    if total_fetched_bytes > max_allowed_bytes {
                        info!(
                            "fetch_blocks() fetched bytes exceeded limit: {} > {}, terminating stream.",
                            total_fetched_bytes, max_allowed_bytes,
                        );
                        break;
                    }
                }
                Ok(None) => {
                    break;
                }
                Err(e) => {
                    if blocks.is_empty() {
                        if e.code() == tonic::Code::DeadlineExceeded {
                            return Err(ConsensusError::NetworkRequestTimeout(format!(
                                "fetch_blocks failed mid-stream: {e:?}"
                            )));
                        }
                        return Err(ConsensusError::NetworkRequest(format!(
                            "fetch_blocks failed mid-stream: {e:?}"
                        )));
                    } else {
                        warn!("fetch_blocks failed mid-stream: {e:?}");
                        break;
                    }
                }
            }
        }
        Ok(blocks)
    }

    async fn fetch_commits(
        &self,
        peer: AuthorityIndex,
        commit_range: CommitRange,
        timeout: Duration,
    ) -> ConsensusResult<(Vec<Bytes>, Vec<Bytes>)> {
        let mut client = self.get_client(peer, timeout).await?;
        let mut request = Request::new(FetchCommitsRequest {
            start: commit_range.start(),
            end: commit_range.end(),
        });
        request.set_timeout(timeout);
        let response = client
            .fetch_commits(request)
            .await
            .map_err(|e| ConsensusError::NetworkRequest(format!("fetch_commits failed: {e:?}")))?;
        let response = response.into_inner();
        Ok((response.commits, response.certifier_blocks))
    }

    async fn fetch_latest_blocks(
        &self,
        peer: AuthorityIndex,
        authorities: Vec<AuthorityIndex>,
        timeout: Duration,
    ) -> ConsensusResult<Vec<Bytes>> {
        let mut client = self.get_client(peer, timeout).await?;
        let mut request = Request::new(FetchLatestBlocksRequest {
            authorities: authorities
                .iter()
                .map(|authority| authority.value() as u32)
                .collect(),
        });
        request.set_timeout(timeout);

        let mut stream = client
            .fetch_latest_blocks(request)
            .await
            .map_err(|e| {
                if e.code() == tonic::Code::DeadlineExceeded {
                    ConsensusError::NetworkRequestTimeout(format!(
                        "fetch_latest_blocks failed: {e:?}"
                    ))
                } else {
                    ConsensusError::NetworkRequest(format!("fetch_latest_blocks failed: {e:?}"))
                }
            })?
            .into_inner();

        // Allow twice the max total size of transactions in the fetched blocks.
        let max_allowed_bytes = authorities.len()
            * self
                .context
                .protocol_config
                .max_transactions_in_block_bytes() as usize
            * 2;
        let mut blocks = vec![];
        let mut total_fetched_bytes = 0;
        loop {
            match stream.message().await {
                Ok(Some(response)) => {
                    for b in &response.blocks {
                        total_fetched_bytes += b.len();
                    }
                    blocks.extend(response.blocks);
                    if total_fetched_bytes > max_allowed_bytes {
                        info!(
                            "fetch_blocks() fetched bytes exceeded limit: {} > {}, terminating stream.",
                            total_fetched_bytes, max_allowed_bytes,
                        );
                        break;
                    }
                }
                Ok(None) => {
                    break;
                }
                Err(e) => {
                    if blocks.is_empty() {
                        if e.code() == tonic::Code::DeadlineExceeded {
                            return Err(ConsensusError::NetworkRequestTimeout(format!(
                                "fetch_blocks failed mid-stream: {e:?}"
                            )));
                        }
                        return Err(ConsensusError::NetworkRequest(format!(
                            "fetch_blocks failed mid-stream: {e:?}"
                        )));
                    } else {
                        warn!("fetch_latest_blocks failed mid-stream: {e:?}");
                        break;
                    }
                }
            }
        }
        Ok(blocks)
    }

    async fn get_latest_rounds(
        &self,
        peer: AuthorityIndex,
        timeout: Duration,
    ) -> ConsensusResult<(Vec<Round>, Vec<Round>)> {
        let mut client = self.get_client(peer, timeout).await?;
        let mut request = Request::new(GetLatestRoundsRequest {});
        request.set_timeout(timeout);
        let response = client.get_latest_rounds(request).await.map_err(|e| {
            ConsensusError::NetworkRequest(format!("get_latest_rounds failed: {e:?}"))
        })?;
        let response = response.into_inner();
        Ok((response.highest_received, response.highest_accepted))
    }

    #[cfg(all(test, not(msim)))]
    async fn send_block(
        &self,
        peer: AuthorityIndex,
        block: &crate::VerifiedBlock,
        timeout: Duration,
    ) -> ConsensusResult<()> {
        let mut client = self.get_client(peer, timeout).await?;
        let mut request = Request::new(SendBlockRequest {
            block: block.serialized().clone(),
        });
        request.set_timeout(timeout);
        client
            .send_block(request)
            .await
            .map_err(|e| ConsensusError::NetworkRequest(format!("send_block failed: {e:?}")))?;
        Ok(())
    }
}

// Tonic channel wrapped with layers.
pub(crate) type Channel = sui_http::middleware::callback::Callback<
    tower::util::MapRequest<
        tower_http::trace::Trace<
            tonic_rustls::Channel,
            tower_http::classify::SharedClassifier<tower_http::classify::GrpcErrorsAsFailures>,
        >,
        ReboxRequestFn,
    >,
    MetricsCallbackMaker,
>;

/// The callback middleware hands the wrapped service a request body of type
/// `RequestBody`, but `tonic_rustls::Channel` is monomorphic on
/// `tonic::body::Body`, so the body must be reboxed before it reaches the
/// channel. A fn pointer (rather than a closure) keeps `Channel` nameable as a
/// type alias.
pub(crate) type ReboxRequestFn =
    fn(http::Request<RequestBody<tonic::body::Body, ()>>) -> http::Request<tonic::body::Body>;

pub(crate) fn rebox_request(
    request: http::Request<RequestBody<tonic::body::Body, ()>>,
) -> http::Request<tonic::body::Body> {
    request.map(tonic::body::Body::new)
}

/// Manages a pool of connections to peers to avoid constantly reconnecting,
/// which can be expensive.
struct ChannelPool {
    context: Arc<Context>,
    // Size is limited by known authorities in the committee.
    channels: RwLock<BTreeMap<AuthorityIndex, Channel>>,
    // Simple address override for peers. If set, this address is used instead of the committee address.
    address_overrides: RwLock<BTreeMap<AuthorityIndex, Multiaddr>>,
}

impl ChannelPool {
    fn new(context: Arc<Context>) -> Self {
        Self {
            context,
            channels: RwLock::new(BTreeMap::new()),
            address_overrides: RwLock::new(BTreeMap::new()),
        }
    }

    // Update the address override for a peer. If None, clears the override.
    fn update_address(&self, peer: AuthorityIndex, address: Option<Multiaddr>) {
        {
            let mut overrides = self.address_overrides.write();

            if let Some(addr) = address {
                if let Some(previous_address) = overrides.insert(peer, addr.clone()) {
                    info!(
                        "Updated address override for peer {}: {} -> {}",
                        peer, previous_address, addr
                    );
                } else {
                    info!("Set address override for peer {} to {}", peer, addr);
                }
            } else if overrides.remove(&peer).is_some() {
                info!("Cleared address override for peer {}", peer);
            }
        }

        // Clear the cached channel so that the next connection attempt uses the updated address.
        let mut channels = self.channels.write();
        if channels.remove(&peer).is_some() {
            info!(
                "Cleared cached channel for peer {} due to address update",
                peer
            );
        }
    }

    async fn get_channel(
        &self,
        network_keypair: NetworkKeyPair,
        peer: AuthorityIndex,
        timeout: Duration,
    ) -> ConsensusResult<Channel> {
        {
            let channels = self.channels.read();
            if let Some(channel) = channels.get(&peer) {
                return Ok(channel.clone());
            }
        }

        let authority = self.context.committee.authority(peer);

        let peer_address = {
            let overrides = self.address_overrides.read();
            overrides
                .get(&peer)
                .cloned()
                .unwrap_or_else(|| authority.address.clone())
        };

        let address = to_host_port_str(&peer_address).map_err(|e| {
            ConsensusError::NetworkConfig(format!("Cannot convert address to host:port: {e:?}"))
        })?;
        let address = format!("https://{address}");
        let config = &self.context.parameters.tonic;
        let buffer_size = config.connection_buffer_size;
        let client_tls_config = sui_tls::create_rustls_client_config(
            self.context
                .committee
                .authority(peer)
                .network_key
                .clone()
                .into_inner(),
            certificate_server_name(&self.context),
            Some(network_keypair.private_key().into_inner()),
        );
        let endpoint = tonic_rustls::Channel::from_shared(address.clone())
            .map_err(|e| ConsensusError::NetworkConfig(format!("invalid URI '{address}': {e}")))?
            .connect_timeout(timeout)
            .initial_connection_window_size(Some(buffer_size as u32))
            .initial_stream_window_size(Some(buffer_size as u32 / 2))
            .keep_alive_while_idle(true)
            .keep_alive_timeout(config.keepalive_interval)
            .http2_keep_alive_interval(config.keepalive_interval)
            // tcp keepalive is probably unnecessary and is unsupported by msim.
            .user_agent("mysticeti")
            .unwrap()
            .tls_config(client_tls_config)
            .unwrap();

        let deadline = tokio::time::Instant::now() + timeout;
        let channel = loop {
            trace!("Connecting to endpoint at {address}");
            match endpoint.connect().await {
                Ok(channel) => break channel,
                Err(e) => {
                    debug!("Failed to connect to endpoint at {address}: {e:?}");
                    if tokio::time::Instant::now() >= deadline {
                        return Err(ConsensusError::NetworkClientConnection(format!(
                            "Timed out connecting to endpoint at {address}: {e:?}"
                        )));
                    }
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
            }
        };
        trace!("Connected to {address}");

        let channel = tower::ServiceBuilder::new()
            .layer(CallbackLayer::new(MetricsCallbackMaker::new(
                self.context.metrics.network_metrics.outbound.clone(),
                self.context.parameters.tonic.excessive_message_size,
            )))
            .map_request(rebox_request as ReboxRequestFn)
            .layer(
                TraceLayer::new_for_grpc()
                    .make_span_with(DefaultMakeSpan::new().level(tracing::Level::TRACE))
                    .on_failure(DefaultOnFailure::new().level(tracing::Level::DEBUG)),
            )
            .service(channel);

        let mut channels = self.channels.write();
        // There should not be many concurrent attempts at connecting to the same peer.
        let channel = channels.entry(peer).or_insert(channel);
        Ok(channel.clone())
    }
}

/// Proxies Tonic requests to NetworkService with actual handler implementation.
struct TonicServiceProxy<S: ValidatorNetworkService> {
    context: Arc<Context>,
    // TonicServiceProxy is cloned into per-connection server tasks, which complete on the
    // network's schedule during graceful shutdown, and can briefly outlive the node if it is
    // dropped without stop(). Hold the service weakly so lingering connections cannot extend
    // the life of the authority service and its state; requests racing shutdown fail with
    // `unavailable` instead.
    service: Weak<S>,
}

impl<S: ValidatorNetworkService> TonicServiceProxy<S> {
    fn new(context: Arc<Context>, service: Arc<S>) -> Self {
        Self {
            context,
            service: Arc::downgrade(&service),
        }
    }

    fn service(&self) -> Result<Arc<S>, tonic::Status> {
        self.service
            .upgrade()
            .ok_or_else(|| tonic::Status::unavailable("Consensus authority is shutting down"))
    }
}

#[async_trait]
impl<S: ValidatorNetworkService> ConsensusService for TonicServiceProxy<S> {
    async fn send_block(
        &self,
        request: Request<SendBlockRequest>,
    ) -> Result<Response<SendBlockResponse>, tonic::Status> {
        let Some(peer_index) = request
            .extensions()
            .get::<PeerInfo>()
            .map(|p| p.authority_index)
        else {
            return Err(tonic::Status::internal("PeerInfo not found"));
        };
        let block = request.into_inner().block;
        // The unicast send path stays raw full SignedBlock bytes; only the live
        // subscription stream uses the envelope framing.
        let block = ExtendedSerializedBlock {
            block: SerializedBlockForm::Full(block),
            excluded_ancestors: vec![],
        };
        self.service()?
            .handle_send_block(peer_index, block)
            .await
            .map_err(|e| tonic::Status::invalid_argument(format!("{e:?}")))?;
        Ok(Response::new(SendBlockResponse {}))
    }

    type SubscribeBlocksStream =
        Pin<Box<dyn Stream<Item = Result<SubscribeBlocksResponse, tonic::Status>> + Send>>;

    async fn subscribe_blocks(
        &self,
        request: Request<Streaming<SubscribeBlocksRequest>>,
    ) -> Result<Response<Self::SubscribeBlocksStream>, tonic::Status> {
        let Some(peer_index) = request
            .extensions()
            .get::<PeerInfo>()
            .map(|p| p.authority_index)
        else {
            return Err(tonic::Status::internal("PeerInfo not found"));
        };
        let mut request_stream = request.into_inner();
        let first_request = match request_stream.next().await {
            Some(Ok(r)) => r,
            Some(Err(e)) => {
                debug!(
                    "subscribe_blocks() request from {} failed: {e:?}",
                    peer_index
                );
                return Err(tonic::Status::invalid_argument("Request error"));
            }
            None => {
                return Err(tonic::Status::invalid_argument("Missing request"));
            }
        };
        let context = self.context.clone();
        let subscriber_hostname = context.committee.authority(peer_index).hostname.clone();
        let stream = self
            .service()?
            .handle_subscribe_blocks(peer_index, first_request.last_received_round)
            .await
            .map_err(|e| tonic::Status::internal(format!("{e:?}")))?
            .map(move |block| {
                let label = match &block.block {
                    SerializedBlockForm::Full(_) => "full",
                    SerializedBlockForm::Slim(_) => "slim",
                };
                let response = SubscribeBlocksResponse {
                    block: frame_subscription_block(
                        block.block,
                        context.protocol_config.slim_block_propagation_enabled(),
                    )?,
                    excluded_ancestors: block.excluded_ancestors,
                };
                // Pre-compression payload bytes by form, so the saving from the
                // slim form is measurable independently of zstd below. With slim
                // propagation enabled the count includes the envelope framing.
                context
                    .metrics
                    .node_metrics
                    .subscribe_blocks_response_bytes
                    .with_label_values(&[subscriber_hostname.as_str(), label])
                    .inc_by(prost::Message::encoded_len(&response) as u64);
                Ok(response)
            });
        let rate_limited_stream =
            tokio_stream::StreamExt::throttle(stream, self.context.parameters.min_round_delay / 2)
                .boxed();
        Ok(Response::new(rate_limited_stream))
    }

    type FetchBlocksStream = Iter<std::vec::IntoIter<Result<FetchBlocksResponse, tonic::Status>>>;

    async fn fetch_blocks(
        &self,
        request: Request<FetchBlocksRequest>,
    ) -> Result<Response<Self::FetchBlocksStream>, tonic::Status> {
        let Some(peer_index) = request
            .extensions()
            .get::<PeerInfo>()
            .map(|p| p.authority_index)
        else {
            return Err(tonic::Status::internal("PeerInfo not found"));
        };
        let inner = request.into_inner();
        let block_refs = inner
            .block_refs
            .into_iter()
            .filter_map(|serialized| match bcs::from_bytes(&serialized) {
                Ok(r) => Some(r),
                Err(e) => {
                    debug!("Failed to deserialize block ref {:?}: {e:?}", serialized);
                    None
                }
            })
            .collect();
        let fetch_after_rounds = inner.fetch_after_rounds;
        let fetch_missing_ancestors = inner.fetch_missing_ancestors;
        let blocks = self
            .service()?
            .handle_fetch_blocks(
                peer_index,
                block_refs,
                fetch_after_rounds,
                fetch_missing_ancestors,
            )
            .await
            .map_err(|e| tonic::Status::internal(format!("{e:?}")))?;
        let responses: std::vec::IntoIter<Result<FetchBlocksResponse, tonic::Status>> =
            chunk_blocks(blocks, MAX_FETCH_RESPONSE_BYTES)
                .into_iter()
                .map(|blocks| Ok(FetchBlocksResponse { blocks }))
                .collect::<Vec<_>>()
                .into_iter();
        let stream = iter(responses);
        Ok(Response::new(stream))
    }

    async fn fetch_commits(
        &self,
        request: Request<FetchCommitsRequest>,
    ) -> Result<Response<FetchCommitsResponse>, tonic::Status> {
        let Some(peer_index) = request
            .extensions()
            .get::<PeerInfo>()
            .map(|p| p.authority_index)
        else {
            return Err(tonic::Status::internal("PeerInfo not found"));
        };
        let request = request.into_inner();
        let (commits, certifier_blocks) = self
            .service()?
            .handle_fetch_commits(peer_index, (request.start..=request.end).into())
            .await
            .map_err(|e| tonic::Status::internal(format!("{e:?}")))?;
        let commits = commits
            .into_iter()
            .map(|c| c.serialized().clone())
            .collect();
        let certifier_blocks = certifier_blocks
            .into_iter()
            .map(|b| b.serialized().clone())
            .collect();
        Ok(Response::new(FetchCommitsResponse {
            commits,
            certifier_blocks,
        }))
    }

    type FetchLatestBlocksStream =
        Iter<std::vec::IntoIter<Result<FetchLatestBlocksResponse, tonic::Status>>>;

    async fn fetch_latest_blocks(
        &self,
        request: Request<FetchLatestBlocksRequest>,
    ) -> Result<Response<Self::FetchLatestBlocksStream>, tonic::Status> {
        let Some(peer_index) = request
            .extensions()
            .get::<PeerInfo>()
            .map(|p| p.authority_index)
        else {
            return Err(tonic::Status::internal("PeerInfo not found"));
        };
        let inner = request.into_inner();

        // Convert the authority indexes and validate them
        let mut authorities = vec![];
        for authority in inner.authorities.into_iter() {
            let Some(authority) = self
                .context
                .committee
                .to_authority_index(authority as usize)
            else {
                return Err(tonic::Status::internal(format!(
                    "Invalid authority index provided {authority}"
                )));
            };
            authorities.push(authority);
        }

        let blocks = self
            .service()?
            .handle_fetch_latest_blocks(peer_index, authorities)
            .await
            .map_err(|e| tonic::Status::internal(format!("{e:?}")))?;
        let responses: std::vec::IntoIter<Result<FetchLatestBlocksResponse, tonic::Status>> =
            chunk_blocks(blocks, MAX_FETCH_RESPONSE_BYTES)
                .into_iter()
                .map(|blocks| Ok(FetchLatestBlocksResponse { blocks }))
                .collect::<Vec<_>>()
                .into_iter();
        let stream = iter(responses);
        Ok(Response::new(stream))
    }

    async fn get_latest_rounds(
        &self,
        request: Request<GetLatestRoundsRequest>,
    ) -> Result<Response<GetLatestRoundsResponse>, tonic::Status> {
        let Some(peer_index) = request
            .extensions()
            .get::<PeerInfo>()
            .map(|p| p.authority_index)
        else {
            return Err(tonic::Status::internal("PeerInfo not found"));
        };
        let (highest_received, highest_accepted) = self
            .service()?
            .handle_get_latest_rounds(peer_index)
            .await
            .map_err(|e| tonic::Status::internal(format!("{e:?}")))?;
        Ok(Response::new(GetLatestRoundsResponse {
            highest_received,
            highest_accepted,
        }))
    }
}

/// Manages the lifecycle of Tonic network client and service. Typical usage during initialization:
/// 1. Create a new `TonicManager`.
/// 2. Take validator and observer clients from `TonicManager::validator_client()` and `TonicManager::observer_client()`.
/// 3. Create consensus components.
/// 4. Create `TonicService` for consensus service handler.
/// 5. Install `TonicService` to `TonicManager` with `TonicManager::install_service()`.
pub(crate) struct TonicManager {
    context: Arc<Context>,
    network_keypair: NetworkKeyPair,
    own_address: SocketAddr,
    validator_client: Arc<TonicValidatorClient>,
    observer_client: Arc<TonicObserverClient>,
    server: Option<ServerHandle>,
    observer_server: Option<ServerHandle>,
}

impl TonicManager {
    pub(crate) fn new(context: Arc<Context>, network_keypair: NetworkKeyPair) -> Self {
        let validator_client = Arc::new(TonicValidatorClient::new(
            context.clone(),
            network_keypair.clone(),
        ));
        let observer_client = Arc::new(TonicObserverClient::new(
            context.clone(),
            network_keypair.clone(),
        ));

        // Calculate own address
        let own_address = if let Some(listen_addr) = &context.parameters.listen_address_override {
            listen_addr
        } else if context.is_observer() {
            // Observer node - use the configured observer server port since we're not in the committee.
            let port = context.parameters.observer.server_port.unwrap_or(0);
            &Multiaddr::try_from(format!("/ip4/0.0.0.0/tcp/{port}")).unwrap()
        } else {
            let authority = context.committee.authority(context.own_index);
            // By default, bind to the unspecified address to allow the actual address to be assigned.
            // But bind to localhost if it is requested.
            if authority.address.is_localhost_ip() {
                &authority.address
            } else {
                &authority.address.with_zero_ip()
            }
        };
        let own_address = to_socket_addr(own_address).unwrap();

        // Calculate TCP metrics for the current machine/OS
        Self::calculate_tcp_metrics(&context, own_address);

        Self {
            context,
            network_keypair,
            own_address,
            validator_client,
            observer_client,
            server: None,
            observer_server: None,
        }
    }

    /// Calculate and record TCP buffer size metrics for the current machine/OS.
    #[cfg(not(msim))]
    fn calculate_tcp_metrics(context: &Arc<Context>, own_address: SocketAddr) {
        let tcp_connection_metrics = &context.metrics.network_metrics.tcp_connection_metrics;

        // Try creating an ephemeral port to test the highest allowed send and recv buffer sizes.
        // Buffer sizes are not set explicitly on the socket used for real traffic,
        // to allow the OS to set appropriate values.
        let ephemeral_addr = SocketAddr::new(own_address.ip(), 0);
        let ephemeral_socket = create_socket(&ephemeral_addr);
        tcp_connection_metrics
            .socket_send_buffer_size
            .set(ephemeral_socket.send_buffer_size().unwrap_or(0) as i64);
        tcp_connection_metrics
            .socket_recv_buffer_size
            .set(ephemeral_socket.recv_buffer_size().unwrap_or(0) as i64);

        if let Err(e) = ephemeral_socket.set_send_buffer_size(32 << 20) {
            info!("Failed to set send buffer size: {e:?}");
        }
        if let Err(e) = ephemeral_socket.set_recv_buffer_size(32 << 20) {
            info!("Failed to set recv buffer size: {e:?}");
        }
        if ephemeral_socket.bind(ephemeral_addr).is_ok() {
            tcp_connection_metrics
                .socket_send_buffer_max_size
                .set(ephemeral_socket.send_buffer_size().unwrap_or(0) as i64);
            tcp_connection_metrics
                .socket_recv_buffer_max_size
                .set(ephemeral_socket.recv_buffer_size().unwrap_or(0) as i64);
        }
    }

    #[cfg(msim)]
    fn calculate_tcp_metrics(_context: &Arc<Context>, _own_address: SocketAddr) {
        // Metrics calculation is not supported in msim
    }
}

impl NetworkManager for TonicManager {
    type ValidatorClient = TonicValidatorClient;
    type ObserverClient = TonicObserverClient;

    fn new(context: Arc<Context>, network_keypair: NetworkKeyPair) -> Self {
        TonicManager::new(context, network_keypair)
    }

    fn validator_client(&self) -> Arc<Self::ValidatorClient> {
        self.validator_client.clone()
    }

    fn observer_client(&self) -> Arc<Self::ObserverClient> {
        self.observer_client.clone()
    }

    fn update_peer_address(&self, peer: AuthorityIndex, address: Option<Multiaddr>) {
        self.validator_client.update_peer_address(peer, address);
    }

    async fn start_validator_server<V>(&mut self, service: Arc<V>)
    where
        V: ValidatorNetworkService,
    {
        info!("Starting tonic validator server");
        self.start_validator_server_impl(service).await;
    }

    async fn start_observer_server<O>(&mut self, service: Arc<O>)
    where
        O: ObserverNetworkService,
    {
        info!("Starting tonic observer server");
        self.start_observer_server_impl(service).await;
    }

    async fn stop(&mut self) {
        if let Some(server) = self.server.take() {
            server.shutdown().await;
        }

        if let Some(observer_server) = self.observer_server.take() {
            observer_server.shutdown().await;
        }
    }
}

impl TonicManager {
    async fn start_validator_server_impl<V: ValidatorNetworkService>(&mut self, service: Arc<V>) {
        let service = TonicServiceProxy::new(self.context.clone(), service);
        let config = &self.context.parameters.tonic;

        let connections_info = Arc::new(ConnectionsInfo::new(self.context.clone()));
        let layers = tower::ServiceBuilder::new()
            // Add a layer to extract a peer's PeerInfo from their TLS certs
            .map_request(move |mut request: http::Request<_>| {
                if let Some(peer_certificates) =
                    request.extensions().get::<sui_http::PeerCertificates>()
                    && let Some(peer_info) =
                        peer_info_from_certs(&connections_info, peer_certificates)
                {
                    request.extensions_mut().insert(peer_info);
                }
                request
            })
            .layer(CallbackLayer::new(MetricsCallbackMaker::new(
                self.context.metrics.network_metrics.inbound.clone(),
                self.context.parameters.tonic.excessive_message_size,
            )))
            .layer(
                TraceLayer::new_for_grpc()
                    .make_span_with(DefaultMakeSpan::new().level(tracing::Level::TRACE))
                    .on_failure(DefaultOnFailure::new().level(tracing::Level::DEBUG)),
            )
            .layer_fn(|service| GrpcTimeout::new(service, Some(DEFAULT_GRPC_SERVER_TIMEOUT)));

        let consensus_service_server = ConsensusServiceServer::new(service)
            .max_encoding_message_size(config.message_size_limit)
            .max_decoding_message_size(config.message_size_limit)
            .send_compressed(CompressionEncoding::Zstd)
            .accept_compressed(CompressionEncoding::Zstd);

        let consensus_service = tonic::service::Routes::new(consensus_service_server)
            .into_axum_router()
            .route_layer(layers);

        let tls_server_config = sui_tls::create_rustls_server_config_with_client_verifier(
            self.network_keypair.clone().private_key().into_inner(),
            certificate_server_name(&self.context),
            AllowPublicKeys::new(
                self.context
                    .committee
                    .authorities()
                    .map(|(_i, a)| a.network_key.clone().into_inner())
                    .collect(),
            ),
        );

        let http_config = sui_http::Config::default()
            .initial_connection_window_size(HTTP2_INITIAL_CONNECTION_WINDOW_SIZE)
            .initial_stream_window_size(HTTP2_INITIAL_STREAM_WINDOW_SIZE)
            .http2_keepalive_interval(Some(config.keepalive_interval))
            .http2_keepalive_timeout(Some(config.keepalive_interval))
            .accept_http1(false);

        // Create server
        //
        // During simtest crash/restart tests there may be an older instance of consensus running
        // that is bound to the TCP port of `own_address` that hasn't finished relinquishing
        // control of the port yet. So instead of crashing when the address is inuse, we will retry
        // for a short/reasonable period of time before giving up.
        let deadline = Instant::now() + Duration::from_secs(20);
        let server = loop {
            match sui_http::Builder::new()
                .config(http_config.clone())
                .tls_config(tls_server_config.clone())
                .serve(self.own_address, consensus_service.clone())
            {
                Ok(server) => break server,
                Err(err) => {
                    warn!("Error starting consensus server: {err:?}");
                    if Instant::now() > deadline {
                        panic!("Failed to start consensus server within required deadline");
                    }
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
            }
        };

        info!("Server started at: {}", self.own_address);
        self.server = Some(server);
    }

    async fn start_observer_server_impl<O: ObserverNetworkService>(&mut self, service: Arc<O>) {
        let observer_params = &self.context.parameters.observer;
        let tonic_config = &self.context.parameters.tonic;
        let Some(observer_port) = observer_params.server_port else {
            info!("Observer server not configured, skipping observer server start");
            return;
        };

        // Parse observer allowlist from configuration and create TLS verifier
        let observer_allowlist = parse_observer_allowlist(&observer_params.allowlist);
        let observer_tls_config = if observer_allowlist.is_empty() {
            info!("Observer server allowlist disabled - all observers allowed");
            sui_tls::create_rustls_server_config_with_client_verifier(
                self.network_keypair.clone().private_key().into_inner(),
                certificate_server_name(&self.context),
                sui_tls::AllowAll,
            )
        } else {
            info!(
                "Observer server allowlist enabled with {} keys",
                observer_allowlist.len()
            );
            let allowed_keys = observer_allowlist
                .into_iter()
                .map(|k| k.into_inner())
                .collect();
            sui_tls::create_rustls_server_config_with_client_verifier(
                self.network_keypair.clone().private_key().into_inner(),
                certificate_server_name(&self.context),
                AllowPublicKeys::new(allowed_keys),
            )
        };

        info!("Starting observer service on port {observer_port}");

        // Use the pre-calculated own address and override with observer port
        let observer_address = SocketAddr::new(self.own_address.ip(), observer_port);
        let observer_service_proxy = ObserverServiceProxy::new(self.context.clone(), service);

        let observer_service_server = ObserverServiceServer::new(observer_service_proxy)
            .max_encoding_message_size(tonic_config.message_size_limit)
            .max_decoding_message_size(tonic_config.message_size_limit)
            .send_compressed(CompressionEncoding::Zstd)
            .accept_compressed(CompressionEncoding::Zstd);

        let layers = tower::ServiceBuilder::new()
            .map_request(move |mut request: http::Request<_>| {
                if let Some(peer_certificates) =
                    request.extensions().get::<sui_http::PeerCertificates>()
                {
                    if let Some(observer_peer_info) =
                        observer_peer_info_from_certs(peer_certificates)
                    {
                        debug!("Inserting observer peer info: {:?}", observer_peer_info);
                        request.extensions_mut().insert(observer_peer_info);
                    }
                } else {
                    debug!("No peer certificates found for observer");
                }
                request
            })
            .layer(CallbackLayer::new(MetricsCallbackMaker::new(
                self.context.metrics.network_metrics.inbound.clone(),
                self.context.parameters.tonic.excessive_message_size,
            )))
            .layer(
                TraceLayer::new_for_grpc()
                    .make_span_with(DefaultMakeSpan::new().level(tracing::Level::TRACE))
                    .on_failure(DefaultOnFailure::new().level(tracing::Level::DEBUG)),
            )
            .layer_fn(|service| GrpcTimeout::new(service, Some(DEFAULT_GRPC_SERVER_TIMEOUT)));

        let observer_service = tonic::service::Routes::new(observer_service_server)
            .into_axum_router()
            .route_layer(layers);

        let http_config = sui_http::Config::default()
            .initial_connection_window_size(HTTP2_INITIAL_CONNECTION_WINDOW_SIZE)
            .initial_stream_window_size(HTTP2_INITIAL_STREAM_WINDOW_SIZE)
            .http2_keepalive_interval(Some(tonic_config.keepalive_interval))
            .http2_keepalive_timeout(Some(tonic_config.keepalive_interval))
            .accept_http1(false);

        let deadline = Instant::now() + Duration::from_secs(20);
        let observer_server = loop {
            match sui_http::Builder::new()
                .config(http_config.clone())
                .tls_config(observer_tls_config.clone())
                .serve(observer_address, observer_service.clone())
            {
                Ok(server) => break server,
                Err(err) => {
                    warn!("Error starting observer server: {err:?}");
                    if Instant::now() > deadline {
                        panic!("Failed to start observer server within required deadline");
                    }
                    tokio::time::sleep(Duration::from_secs(1)).await;
                }
            }
        };

        info!("Observer server started at: {observer_address}");
        self.observer_server = Some(observer_server);
    }
}

// Ensure that if there is an active network running that it is shutdown when the TonicManager is
// dropped.
impl Drop for TonicManager {
    fn drop(&mut self) {
        if let Some(server) = self.server.as_ref() {
            server.trigger_shutdown();
        }
        if let Some(observer_server) = self.observer_server.as_ref() {
            observer_server.trigger_shutdown();
        }
    }
}

// TODO: improve sui-http to allow for providing a MakeService so that this can be done once per
// connection
fn peer_info_from_certs(
    connections_info: &ConnectionsInfo,
    peer_certificates: &sui_http::PeerCertificates,
) -> Option<PeerInfo> {
    let certs = peer_certificates.peer_certs();

    if certs.len() != 1 {
        trace!(
            "Unexpected number of certificates from TLS stream: {}",
            certs.len()
        );
        return None;
    }
    trace!("Received {} certificates", certs.len());
    let public_key = sui_tls::public_key_from_certificate(&certs[0])
        .map_err(|e| {
            trace!("Failed to extract public key from certificate: {e:?}");
            e
        })
        .ok()?;
    let client_public_key = NetworkPublicKey::new(public_key);
    let Some(authority_index) = connections_info.authority_index(&client_public_key) else {
        error!("Failed to find the authority with public key {client_public_key:?}");
        return None;
    };
    Some(PeerInfo { authority_index })
}

/// Extracts observer peer information from TLS certificates.
/// Unlike validator peers, observers are not required to be in the committee.
/// The allowlist filtering is enforced at the TLS level via AllowPublicKeys or AllowAll.
fn observer_peer_info_from_certs(
    peer_certificates: &sui_http::PeerCertificates,
) -> Option<ObserverPeerInfo> {
    let certs = peer_certificates.peer_certs();

    if certs.len() != 1 {
        trace!(
            "Unexpected number of certificates from TLS stream: {}",
            certs.len()
        );
        return None;
    }
    trace!("Received {} observer certificates", certs.len());
    let public_key = sui_tls::public_key_from_certificate(&certs[0])
        .map_err(|e| {
            trace!("Failed to extract public key from observer certificate: {e:?}");
            e
        })
        .ok()?;
    let client_public_key = NetworkPublicKey::new(public_key);

    debug!(
        "Observer connection accepted: public key {:?}",
        client_public_key
    );

    Some(ObserverPeerInfo {
        public_key: client_public_key,
    })
}

/// Parses the observer allowlist from hex-encoded strings to NetworkPublicKey objects.
/// Returns an empty Vec if the allowlist is empty (meaning all observers are allowed).
/// Logs errors for any invalid keys and skips them.
fn parse_observer_allowlist(allowlist_strings: &[String]) -> Vec<NetworkPublicKey> {
    if allowlist_strings.is_empty() {
        return Vec::new();
    }

    allowlist_strings
        .iter()
        .filter_map(|key_str| match fastcrypto::encoding::Hex::decode(key_str) {
            Ok(bytes) => match fastcrypto::ed25519::Ed25519PublicKey::from_bytes(bytes.as_ref()) {
                Ok(inner_key) => Some(NetworkPublicKey::new(inner_key)),
                Err(e) => {
                    error!(
                        "Failed to parse observer public key from bytes '{}': {e:?}",
                        key_str
                    );
                    None
                }
            },
            Err(e) => {
                error!(
                    "Failed to decode hex observer public key '{}': {e:?}",
                    key_str
                );
                None
            }
        })
        .collect()
}

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

    match (iter.next(), iter.next()) {
        (Some(Protocol::Ip4(ipaddr)), Some(Protocol::Udp(port)))
        | (Some(Protocol::Ip4(ipaddr)), Some(Protocol::Tcp(port))) => {
            Ok(SocketAddr::V4(SocketAddrV4::new(ipaddr, port)))
        }

        (Some(Protocol::Ip6(ipaddr)), Some(Protocol::Udp(port)))
        | (Some(Protocol::Ip6(ipaddr)), Some(Protocol::Tcp(port))) => {
            Ok(SocketAddr::V6(SocketAddrV6::new(ipaddr, port, 0, 0)))
        }

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

#[cfg(not(msim))]
fn create_socket(address: &SocketAddr) -> tokio::net::TcpSocket {
    let socket = if address.is_ipv4() {
        tokio::net::TcpSocket::new_v4()
    } else if address.is_ipv6() {
        tokio::net::TcpSocket::new_v6()
    } else {
        panic!("Invalid own address: {address:?}");
    }
    .unwrap_or_else(|e| panic!("Cannot create TCP socket: {e:?}"));
    if let Err(e) = socket.set_nodelay(true) {
        info!("Failed to set TCP_NODELAY: {e:?}");
    }
    if let Err(e) = socket.set_reuseaddr(true) {
        info!("Failed to set SO_REUSEADDR: {e:?}");
    }
    socket
}

/// Looks up authority index by authority public key.
///
/// TODO: Add connection monitoring, and keep track of connected peers.
/// TODO: Maybe merge with connection_monitor.rs
struct ConnectionsInfo {
    authority_key_to_index: BTreeMap<NetworkPublicKey, AuthorityIndex>,
}

impl ConnectionsInfo {
    fn new(context: Arc<Context>) -> Self {
        let authority_key_to_index = context
            .committee
            .authorities()
            .map(|(index, authority)| (authority.network_key.clone(), index))
            .collect();
        Self {
            authority_key_to_index,
        }
    }

    fn authority_index(&self, key: &NetworkPublicKey) -> Option<AuthorityIndex> {
        self.authority_key_to_index.get(key).copied()
    }
}

/// Information about the client peer, set per connection.
#[derive(Clone, Debug)]
struct PeerInfo {
    authority_index: AuthorityIndex,
}

// Adapt MetricsCallbackMaker and MetricsResponseCallback to http.

impl SizedRequest for http::request::Parts {
    fn size(&self) -> usize {
        // TODO: implement this.
        0
    }

    fn route(&self) -> String {
        let path = self.uri.path();
        path.rsplit_once('/')
            .map(|(_, route)| route)
            .unwrap_or("unknown")
            .to_string()
    }
}

impl SizedResponse for http::response::Parts {
    fn size(&self) -> usize {
        // TODO: implement this.
        0
    }

    fn error_type(&self) -> Option<String> {
        if self.status.is_success() {
            None
        } else {
            Some(self.status.to_string())
        }
    }
}

impl MakeCallbackHandler for MetricsCallbackMaker {
    type RequestHandler = ();
    type ResponseHandler = MetricsResponseCallback;

    fn make_handler(
        &self,
        request: &http::request::Parts,
    ) -> (Self::RequestHandler, Self::ResponseHandler) {
        ((), self.handle_request(request))
    }
}

impl ResponseHandler for MetricsResponseCallback {
    fn on_response(&mut self, response: &http::response::Parts) {
        MetricsResponseCallback::on_response(self, response)
    }

    fn on_service_error<E>(&mut self, err: &E)
    where
        E: std::fmt::Display + 'static,
    {
        MetricsResponseCallback::on_error(self, err)
    }
}

/// Network message types.
#[derive(Clone, prost::Message)]
pub(crate) struct SendBlockRequest {
    // Serialized SignedBlock.
    #[prost(bytes = "bytes", tag = "1")]
    block: Bytes,
}

#[derive(Clone, prost::Message)]
pub(crate) struct SendBlockResponse {}

#[derive(Clone, prost::Message)]
pub(crate) struct SubscribeBlocksRequest {
    #[prost(uint32, tag = "1")]
    last_received_round: Round,
}

/// Frames one block for the live subscription wire. With slim propagation enabled every
/// payload rides in a [`SerializedBlockEnvelope`], whichever form it is; disabled, the
/// bytes are the raw serialized `SignedBlock`, byte-identical to the pre-envelope wire.
/// A slim form while the flag is off is an internal invariant violation -- nothing
/// produces one -- and fails the stream rather than fabricating a payload the receiver
/// would reject block by block.
fn frame_subscription_block(
    form: SerializedBlockForm,
    slim_enabled: bool,
) -> Result<Bytes, tonic::Status> {
    if slim_enabled {
        return Ok(SerializedBlockEnvelope::encode_form(form));
    }
    match form {
        SerializedBlockForm::Full(bytes) => Ok(bytes),
        SerializedBlockForm::Slim(_) => Err(tonic::Status::internal(
            "slim block emitted while slim propagation is disabled",
        )),
    }
}

/// Inverse of [`frame_subscription_block`], on the receiving side of the stream.
fn interpret_subscription_block(
    block: Bytes,
    slim_enabled: bool,
) -> ConsensusResult<SerializedBlockForm> {
    if slim_enabled {
        SerializedBlockEnvelope::decode_form(&block)
    } else {
        Ok(SerializedBlockForm::Full(block))
    }
}

#[derive(Clone, prost::Message)]
pub(crate) struct SubscribeBlocksResponse {
    // With slim block propagation disabled these are raw serialized SignedBlock bytes,
    // byte-identical to the pre-envelope wire. With it enabled they hold a serialized
    // SerializedBlockEnvelope.
    #[prost(bytes = "bytes", tag = "1")]
    block: Bytes,
    // Serialized BlockRefs that are excluded from the blocks ancestors.
    #[prost(bytes = "vec", repeated, tag = "2")]
    excluded_ancestors: Vec<Vec<u8>>,
}

#[derive(Clone, prost::Message)]
pub(crate) struct FetchBlocksRequest {
    #[prost(bytes = "vec", repeated, tag = "1")]
    block_refs: Vec<Vec<u8>>,
    // The round per authority after which blocks should be fetched. The vector represents the round
    // for each authority and its length should be the same as the committee size.
    // When this field is non-empty, additional ancestors of the requested blocks can be fetched.
    #[prost(uint32, repeated, tag = "2")]
    fetch_after_rounds: Vec<Round>,
    // When true, missing ancestors of the requested blocks will be fetched as well.
    // When false, additional blocks are fetched depth-first from the requested block authorities.
    // This field is only meaningful when fetch_after_rounds is non-empty.
    #[prost(bool, tag = "3")]
    fetch_missing_ancestors: bool,
}

#[derive(Clone, prost::Message)]
pub(crate) struct FetchBlocksResponse {
    // The response of the requested blocks as Serialized SignedBlock.
    #[prost(bytes = "bytes", repeated, tag = "1")]
    blocks: Vec<Bytes>,
}

#[derive(Clone, prost::Message)]
pub(crate) struct FetchCommitsRequest {
    #[prost(uint32, tag = "1")]
    start: CommitIndex,
    #[prost(uint32, tag = "2")]
    end: CommitIndex,
}

#[derive(Clone, prost::Message)]
pub(crate) struct FetchCommitsResponse {
    // Serialized consecutive Commit.
    #[prost(bytes = "bytes", repeated, tag = "1")]
    commits: Vec<Bytes>,
    // Serialized SignedBlock that certify the last commit from above.
    #[prost(bytes = "bytes", repeated, tag = "2")]
    certifier_blocks: Vec<Bytes>,
}

#[derive(Clone, prost::Message)]
pub(crate) struct FetchLatestBlocksRequest {
    #[prost(uint32, repeated, tag = "1")]
    authorities: Vec<u32>,
}

#[derive(Clone, prost::Message)]
pub(crate) struct FetchLatestBlocksResponse {
    // The response of the requested blocks as Serialized SignedBlock.
    #[prost(bytes = "bytes", repeated, tag = "1")]
    blocks: Vec<Bytes>,
}

#[derive(Clone, prost::Message)]
pub(crate) struct GetLatestRoundsRequest {}

#[derive(Clone, prost::Message)]
pub(crate) struct GetLatestRoundsResponse {
    // Highest received round per authority.
    #[prost(uint32, repeated, tag = "1")]
    highest_received: Vec<u32>,
    // Highest accepted round per authority.
    #[prost(uint32, repeated, tag = "2")]
    highest_accepted: Vec<u32>,
}

pub(crate) fn chunk_blocks(blocks: Vec<Bytes>, chunk_limit: usize) -> Vec<Vec<Bytes>> {
    let mut chunks = vec![];
    let mut chunk = vec![];
    let mut chunk_size = 0;
    for block in blocks {
        let block_size = block.len();
        if !chunk.is_empty() && chunk_size + block_size > chunk_limit {
            chunks.push(chunk);
            chunk = vec![];
            chunk_size = 0;
        }
        chunk.push(block);
        chunk_size += block_size;
    }
    if !chunk.is_empty() {
        chunks.push(chunk);
    }
    chunks
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{context::Clock, metrics::initialise_metrics};
    use consensus_config::{ConsensusProtocolConfig, Parameters, local_committee_and_keys};
    use consensus_types::block::BlockDigest;
    use prometheus::Registry;

    /// With slim propagation disabled the framing is the identity: the wire carries
    /// the raw SignedBlock bytes unchanged, in the same message struct as before the
    /// envelope existed -- so deploying this is a wire no-op.
    #[tokio::test]
    async fn flag_off_wire_carries_raw_block_bytes() {
        for block in [Bytes::from_static(b"a-signed-block"), Bytes::new()] {
            assert_eq!(
                frame_subscription_block(SerializedBlockForm::Full(block.clone()), false).unwrap(),
                block
            );
        }
    }

    /// With the flag on, both forms survive framing and interpretation.
    #[tokio::test]
    async fn flag_on_roundtrips_both_forms() {
        for form in [
            SerializedBlockForm::Full(Bytes::from_static(b"full-bytes")),
            SerializedBlockForm::Slim(Bytes::from_static(b"slim-bytes")),
        ] {
            let framed = frame_subscription_block(form.clone(), true).unwrap();
            assert_eq!(interpret_subscription_block(framed, true).unwrap(), form);
        }
        // Empty payloads are still valid envelopes, distinct from an absent form.
        for form in [
            SerializedBlockForm::Full(Bytes::new()),
            SerializedBlockForm::Slim(Bytes::new()),
        ] {
            let framed = frame_subscription_block(form.clone(), true).unwrap();
            assert_eq!(interpret_subscription_block(framed, true).unwrap(), form);
        }
    }

    /// A slim form with the flag off is an internal invariant violation and must fail
    /// the stream, not fabricate a payload.
    #[tokio::test]
    async fn flag_off_slim_fails_the_stream() {
        assert!(
            frame_subscription_block(
                SerializedBlockForm::Slim(Bytes::from_static(b"slim")),
                false
            )
            .is_err()
        );
    }

    /// The framing follows the protocol config end to end: the testing setter flips the
    /// flag on a real config, and both framing directions change with it.
    #[tokio::test]
    async fn framing_follows_the_protocol_config() {
        let mut config = consensus_config::ConsensusProtocolConfig::for_testing();
        assert!(
            !config.slim_block_propagation_enabled(),
            "the flag must default off, including for tests"
        );
        let raw = Bytes::from_static(b"a-signed-block");

        let framed_off = frame_subscription_block(
            SerializedBlockForm::Full(raw.clone()),
            config.slim_block_propagation_enabled(),
        )
        .unwrap();
        assert_eq!(framed_off, raw, "flag off leaves the bytes untouched");

        config.set_slim_block_propagation_enabled_for_testing(true);
        let framed_on = frame_subscription_block(
            SerializedBlockForm::Full(raw.clone()),
            config.slim_block_propagation_enabled(),
        )
        .unwrap();
        assert_ne!(framed_on, raw, "flag on envelopes the bytes");
        assert_eq!(
            interpret_subscription_block(framed_on, config.slim_block_propagation_enabled())
                .unwrap(),
            SerializedBlockForm::Full(raw)
        );
    }

    /// Envelope decode failures a receiver must survive: garbage, an empty envelope,
    /// and raw BCS block bytes sent by a flag-off (or lying) peer. Raw SignedBlock
    /// bytes always fail: their first byte is the block-version index (1..=3), which
    /// protobuf reads as a tag for the illegal field number 0.
    #[tokio::test]
    async fn malformed_envelopes_are_errors_not_blocks() {
        assert!(matches!(
            interpret_subscription_block(Bytes::from_static(b"\xff\xff\xff"), true),
            Err(ConsensusError::MalformedBlockEnvelope(_))
        ));
        assert!(
            matches!(
                interpret_subscription_block(Bytes::new(), true),
                Err(ConsensusError::MalformedBlockEnvelope(_))
            ),
            "an empty envelope has no form and must not decode"
        );
        for version_byte in [1u8, 2, 3] {
            let raw = Bytes::from(vec![version_byte, 0x42, 0x42, 0x42]);
            assert!(
                interpret_subscription_block(raw, true).is_err(),
                "raw BCS SignedBlock bytes must not parse as an envelope"
            );
        }
    }

    /// The unsupported mismatch direction the other way: a flag-off receiver handed an
    /// enveloped payload treats it as raw block bytes, which then fail BCS
    /// deserialization downstream -- rejected per block, never accepted.
    #[tokio::test]
    async fn flag_off_receiver_passes_enveloped_bytes_through() {
        let enveloped =
            frame_subscription_block(SerializedBlockForm::Full(Bytes::from_static(b"x")), true)
                .unwrap();
        let form = interpret_subscription_block(enveloped.clone(), false).unwrap();
        assert_eq!(form, SerializedBlockForm::Full(enveloped));
    }

    fn create_test_context_and_client() -> (Arc<Context>, TonicValidatorClient) {
        let (committee, mut keypairs) = local_committee_and_keys(0, vec![1, 1, 1, 1]);
        let parameters = Parameters::default();
        let protocol_config = ConsensusProtocolConfig::for_testing();
        let metrics = initialise_metrics(Registry::new());

        let authority_index = committee.to_authority_index(0).unwrap();
        let context = Arc::new(Context::new(
            0,
            Some(authority_index),
            committee,
            parameters,
            protocol_config,
            metrics,
            Arc::new(Clock::default()),
        ));

        let (network_keypair, _protocol_keypair) = keypairs.remove(0);
        let client = TonicValidatorClient::new(context.clone(), network_keypair);

        (context, client)
    }

    #[tokio::test]
    async fn test_max_fetch_blocks_response_bytes_uses_response_mode_limit() {
        let (mut context, _keys) = Context::new_for_test(4);
        context.parameters.max_blocks_per_fetch = 11;
        context.parameters.max_blocks_per_sync = 7;
        let context = Arc::new(context);
        let bytes_per_block =
            context.protocol_config.max_transactions_in_block_bytes() as usize * 2;

        let block_ref = BlockRef::new(1, AuthorityIndex::new_for_test(0), BlockDigest::MIN);
        let fetch_after_rounds = vec![0; context.committee.size()];

        // Commit-sync mode (block_refs only): always sized to max_blocks_per_fetch,
        // independent of how many block_refs the caller actually requested.
        assert_eq!(
            max_fetch_blocks_response_bytes(&context, &[block_ref; 3], &[]),
            context.parameters.max_blocks_per_fetch * bytes_per_block
        );
        assert_eq!(
            max_fetch_blocks_response_bytes(&context, &[block_ref; 20], &[]),
            context.parameters.max_blocks_per_fetch * bytes_per_block
        );
        // Periodic-sync mode (fetch_after_rounds only): max_blocks_per_fetch.
        assert_eq!(
            max_fetch_blocks_response_bytes(&context, &[], &fetch_after_rounds),
            context.parameters.max_blocks_per_fetch * bytes_per_block
        );
        // Live-sync mode (both inputs non-empty): max_blocks_per_sync.
        assert_eq!(
            max_fetch_blocks_response_bytes(&context, &[block_ref], &fetch_after_rounds),
            context.parameters.max_blocks_per_sync * bytes_per_block
        );
    }

    #[tokio::test]
    async fn test_update_peer_address() {
        let (context, client) = create_test_context_and_client();

        let peer = context.committee.to_authority_index(1).unwrap();
        let new_address: Multiaddr = "/ip4/127.0.0.1/udp/9000".parse().unwrap();

        client.update_peer_address(peer, Some(new_address.clone()));

        // Verify the override was set
        {
            let overrides = client.channel_pool.address_overrides.read();
            assert_eq!(overrides.get(&peer), Some(&new_address));
        }

        // Update with a different address
        let newer_address: Multiaddr = "/ip4/127.0.0.1/udp/9001".parse().unwrap();
        client.update_peer_address(peer, Some(newer_address.clone()));

        // Verify the override was updated
        {
            let overrides = client.channel_pool.address_overrides.read();
            assert_eq!(overrides.get(&peer), Some(&newer_address));
        }

        // Verify channels map doesn't contain the peer
        {
            let channels = client.channel_pool.channels.read();
            assert!(!channels.contains_key(&peer));
        }
    }

    #[tokio::test]
    async fn test_clear_peer_address() {
        let (context, client) = create_test_context_and_client();

        let peer = context.committee.to_authority_index(1).unwrap();
        let new_address: Multiaddr = "/ip4/127.0.0.1/udp/9000".parse().unwrap();

        // Set address override
        client.update_peer_address(peer, Some(new_address));

        // Verify the override was set
        {
            let overrides = client.channel_pool.address_overrides.read();
            assert!(overrides.contains_key(&peer));
        }

        // Clear the override
        client.update_peer_address(peer, None);

        // Verify the override was cleared
        {
            let overrides = client.channel_pool.address_overrides.read();
            assert!(!overrides.contains_key(&peer));
        }
    }

    #[tokio::test]
    async fn test_different_peers_independent() {
        let (context, client) = create_test_context_and_client();

        let peer1 = context.committee.to_authority_index(1).unwrap();
        let peer2 = context.committee.to_authority_index(2).unwrap();

        let address1: Multiaddr = "/ip4/127.0.0.1/udp/9000".parse().unwrap();
        let address2: Multiaddr = "/ip4/127.0.0.1/udp/9001".parse().unwrap();

        // Set different addresses for different peers
        client.update_peer_address(peer1, Some(address1.clone()));
        client.update_peer_address(peer2, Some(address2.clone()));

        // Verify both overrides are set correctly
        {
            let overrides = client.channel_pool.address_overrides.read();
            assert_eq!(overrides.get(&peer1), Some(&address1));
            assert_eq!(overrides.get(&peer2), Some(&address2));
        }

        // Clear one peer's override
        client.update_peer_address(peer1, None);

        // Verify only peer1's override was cleared
        {
            let overrides = client.channel_pool.address_overrides.read();
            assert!(!overrides.contains_key(&peer1));
            assert_eq!(overrides.get(&peer2), Some(&address2));
        }
    }
}