//! Messages and bytes exchanged with each peer, by kind: counted where every message
//! passes anyway, at the transport. No messages of our own.
//!
//! Bytes are payload bytes, what a message carries, without the transport's framing
//! (for Mycelium: the topic, or the base64 of its HTTP API). Counters only grow (from
//! the start of the process); rates are for whoever reads them.
//!
//! A message dropped by the debug firewall (in the paths layer, below) still counts as
//! sent (the firewall plays a lossy network) but not as received.

use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};

use anyhow::Result;
use async_trait::async_trait;
use serde::Serialize;

use crate::locks::MutexExt;
use crate::transport::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};

#[derive(Debug, Default, Clone, Copy, Serialize)]
pub(crate) struct Counter {
    pub(crate) messages: u64,
    pub(crate) bytes: u64,
}

impl Counter {
    fn add(&mut self, bytes: usize) {
        self.messages += 1;
        self.bytes += bytes as u64;
    }
}

#[derive(Debug, Default, Clone, Copy, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct ByKind {
    pub(crate) hello: Counter,
    pub(crate) hello_reply: Counter,
    pub(crate) sync: Counter,
    pub(crate) status: Counter,
    pub(crate) status_reply: Counter,
    pub(crate) blob: Counter,
    /// WHO and ME (`transport::identity`), counted by the identity layer under the
    /// address they went to or came from.
    pub(crate) identity: Counter,
    /// PING and PONG (`transport::paths`), counted by the paths layer under the address.
    pub(crate) paths: Counter,
    /// LINKS and LINKS_REPLY (`swarm::links`).
    pub(crate) links: Counter,
}

impl ByKind {
    fn get_mut(&mut self, kind: MessageKind) -> &mut Counter {
        match kind {
            MessageKind::Hello => &mut self.hello,
            MessageKind::HelloReply => &mut self.hello_reply,
            MessageKind::Sync => &mut self.sync,
            MessageKind::Status => &mut self.status,
            MessageKind::StatusReply => &mut self.status_reply,
            MessageKind::Blob => &mut self.blob,
            MessageKind::Who | MessageKind::Me => &mut self.identity,
            MessageKind::Ping | MessageKind::Pong => &mut self.paths,
            MessageKind::Links | MessageKind::LinksReply => &mut self.links,
        }
    }
}

#[derive(Debug, Default, Clone, Copy, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PeerTraffic {
    /// Handed to the transport (it accepted them).
    pub(crate) sent: ByKind,
    pub(crate) received: ByKind,
    /// Sends the transport refused (e.g. no route to the peer); not counted as sent.
    pub(crate) refused: ByKind,
}

#[derive(Debug, Default)]
struct Table {
    peers: HashMap<String, PeerTraffic>,
    /// All peers, including ones no longer listed.
    total: PeerTraffic,
}

#[derive(Debug, Default)]
pub(crate) struct Traffic {
    table: Mutex<Table>,
    /// The messages one by one, while the dashboard looks.
    pub(crate) log: MessageLog,
}

/// Messages kept in the log (`MessageLog`) at most, and the text of one.
const LOGGED: usize = 1000;
const LOGGED_TEXT: usize = 8 * 1024;

/// The node's messages one by one, for the dashboard's console (`api::dashboard`):
/// kept only while it looks (`watch`), the latest `LOGGED`. Those this layer sees: the
/// node's own, not the paths layer's PINGs and PONGs nor the identity layer's WHO and
/// ME (counted under their addresses, `ByKind`), but the PINGs and PONGs the node is
/// handed.
#[derive(Debug, Default)]
pub(crate) struct MessageLog {
    watching: AtomicBool,
    inner: Mutex<LogInner>,
}

#[derive(Debug, Default)]
struct LogInner {
    next: u64,
    entries: VecDeque<Logged>,
}

/// One message, as the dashboard shows it.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct Logged {
    /// Counts up, from 1.
    pub(crate) seq: u64,
    /// Unix milliseconds, by the system's clock.
    pub(crate) at: u64,
    /// `sent`, `received` or `refused` (the transport wouldn't send it).
    pub(crate) way: &'static str,
    pub(crate) peer: String,
    /// The message's type, as the docs name it: `HELLO_REPLY`, `SYNC`, …
    #[serde(rename = "type")]
    pub(crate) kind: &'static str,
    /// Its kind among the counters (`ByKind`'s names): `helloReply`, `sync`, …
    pub(crate) counted: &'static str,
    pub(crate) bytes: usize,
    /// The message, written out (`Debug`; sizes for its bytes), `LOGGED_TEXT` at most.
    pub(crate) text: String,
    /// What ties it to other messages (`related`): those with the same are one
    /// exchange.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) related: Option<String>,
    /// The address it went to or came from, where the paths layer logged it under that
    /// (its own PINGs and PONGs) and the node there is known: `peer` is the node then
    /// (`api::dashboard`).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) address: Option<String>,
}

/// What ties `message`, with `peer`, to other messages: a PING and its PONG (the
/// nonce alone: a PING may go to an address and its PONG come back up from the node
/// there, `paths`), a LINKS and its replies (the nonce), a HELLO and its replies (the run of the
/// node that sent the HELLO, which the reply echoes), a blob's requests and answers, from
/// any peer (its ID), and a delta offered in a SYNC with the blob it is fetched as. None
/// for the rest (a STATUS and its reply carry nothing of the kind).
fn related(peer: &str, message: &ParsedMessage) -> Option<String> {
    use crate::transport::{Hex, SyncBody};
    Some(match message {
        ParsedMessage::Ping(nonce) | ParsedMessage::Pong { nonce, .. } => {
            format!("ping {nonce}")
        }
        ParsedMessage::Links(request) => format!("links {peer} {}", request.nonce),
        ParsedMessage::LinksReply(reply) => format!("links {peer} {}", reply.nonce),
        ParsedMessage::Hello(hello) => format!("hello {peer} {}", Hex(&hello.run_id)),
        ParsedMessage::HelloReply(reply) => format!("hello {peer} {}", Hex(&reply.run_id)),
        ParsedMessage::Blob(blob) => format!("blob {}", blob.id),
        ParsedMessage::Sync(envelope) => match &envelope.body {
            SyncBody::Delta(spec) => {
                format!("blob {}", crate::transport::BlobId::Delta(spec.id()))
            }
            SyncBody::Message(_) => return None,
        },
        _ => return None,
    })
}

// Read by the dashboard only.
#[cfg_attr(not(feature = "dashboard"), allow(dead_code))]
impl MessageLog {
    /// Keeps messages from now on (`true`), or no more, forgetting those kept.
    pub(crate) fn watch(&self, on: bool) {
        self.watching.store(on, Ordering::Relaxed);
        if !on {
            self.inner.lock_or_recover().entries.clear();
        }
    }

    pub(crate) fn note(
        &self,
        way: &'static str,
        peer: &str,
        message: &ParsedMessage,
        bytes: usize,
    ) {
        if !self.watching.load(Ordering::Relaxed) {
            return;
        }
        let at = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |since| since.as_millis() as u64);
        let mut text = format!("{message:#?}");
        if text.len() > LOGGED_TEXT {
            let mut cut = LOGGED_TEXT;
            while !text.is_char_boundary(cut) {
                cut -= 1;
            }
            text.truncate(cut);
            text.push_str("\n…");
        }
        let (kind, counted) = names(message.kind());
        let mut inner = self.inner.lock_or_recover();
        inner.next += 1;
        let seq = inner.next;
        inner.entries.push_back(Logged {
            seq,
            at,
            way,
            peer: peer.to_owned(),
            kind,
            counted,
            bytes,
            text,
            related: related(peer, message),
            address: None,
        });
        if inner.entries.len() > LOGGED {
            inner.entries.pop_front();
        }
    }

    /// The messages after `seq`, the latest `most` of them.
    pub(crate) fn after(&self, seq: u64, most: usize) -> Vec<Logged> {
        let inner = self.inner.lock_or_recover();
        let newer: Vec<&Logged> = inner.entries.iter().filter(|e| e.seq > seq).collect();
        newer[newer.len().saturating_sub(most)..]
            .iter()
            .map(|entry| (*entry).clone())
            .collect()
    }
}

/// A message kind's type name, and its counter's (`ByKind`).
fn names(kind: MessageKind) -> (&'static str, &'static str) {
    match kind {
        MessageKind::Hello => ("HELLO", "hello"),
        MessageKind::HelloReply => ("HELLO_REPLY", "helloReply"),
        MessageKind::Sync => ("SYNC", "sync"),
        MessageKind::Status => ("STATUS", "status"),
        MessageKind::StatusReply => ("STATUS_REPLY", "statusReply"),
        MessageKind::Blob => ("BLOB", "blob"),
        MessageKind::Who => ("WHO", "identity"),
        MessageKind::Me => ("ME", "identity"),
        MessageKind::Ping => ("PING", "paths"),
        MessageKind::Pong => ("PONG", "paths"),
        MessageKind::Links => ("LINKS", "links"),
        MessageKind::LinksReply => ("LINKS_REPLY", "links"),
    }
}

impl Traffic {
    fn lock(&self) -> MutexGuard<'_, Table> {
        self.table.lock_or_recover()
    }

    fn count(&self, peer: &str, bytes: usize, pick: impl Fn(&mut PeerTraffic) -> &mut Counter) {
        let mut guard = self.lock();
        let table = &mut *guard;
        pick(&mut table.total).add(bytes);
        match table.peers.get_mut(peer) {
            Some(traffic) => pick(traffic).add(bytes),
            None => pick(table.peers.entry(peer.to_owned()).or_default()).add(bytes),
        }
    }

    pub(crate) fn sent(&self, peer: &str, kind: MessageKind, bytes: usize) {
        self.count(peer, bytes, |t| t.sent.get_mut(kind));
    }

    pub(crate) fn refused(&self, peer: &str, kind: MessageKind, bytes: usize) {
        self.count(peer, bytes, |t| t.refused.get_mut(kind));
    }

    pub(crate) fn received(&self, peer: &str, kind: MessageKind, bytes: usize) {
        self.count(peer, bytes, |t| t.received.get_mut(kind));
    }

    pub(crate) fn peer(&self, peer: &str) -> Option<PeerTraffic> {
        self.lock().peers.get(peer).copied()
    }

    pub(crate) fn total(&self) -> PeerTraffic {
        self.lock().total
    }

    /// Drops the counters of peers not kept (anyone can send us messages; the totals
    /// keep their share).
    pub(crate) fn retain(&self, mut keep: impl FnMut(&str) -> bool) {
        self.lock().peers.retain(|peer, _| keep(peer));
    }
}

/// Counts what passes through `inner`.
pub(crate) struct CountingTransport {
    inner: Arc<dyn Transport>,
    traffic: Arc<Traffic>,
}

impl CountingTransport {
    pub(crate) fn new(inner: Arc<dyn Transport>, traffic: Arc<Traffic>) -> Self {
        Self { inner, traffic }
    }
}

#[async_trait]
impl Transport for CountingTransport {
    fn local_id(&self) -> &str {
        self.inner.local_id()
    }

    fn is_peer_id(&self, id: &str) -> bool {
        self.inner.is_peer_id(id)
    }

    fn delivery_limit(&self) -> std::time::Duration {
        self.inner.delivery_limit()
    }

    fn sync_reply_timeout(&self) -> std::time::Duration {
        self.inner.sync_reply_timeout()
    }

    fn local_addresses(&self) -> Vec<String> {
        self.inner.local_addresses()
    }

    fn addresses_of(&self, peer: &str) -> Vec<String> {
        self.inner.addresses_of(peer)
    }

    fn identities_of(&self, peer: &str) -> Vec<String> {
        self.inner.identities_of(peer)
    }

    fn learn_addresses(&self, peer: &str, addresses: &[String]) {
        self.inner.learn_addresses(peer, addresses)
    }

    fn can_reach(&self, peer: &str) -> bool {
        self.inner.can_reach(peer)
    }

    fn discovered_seeds(&self) -> Vec<String> {
        self.inner.discovered_seeds()
    }

    fn on_seeds_found(&self, found: std::sync::Arc<tokio::sync::Notify>) {
        self.inner.on_seeds_found(found)
    }

    fn transports(&self) -> Vec<String> {
        self.inner.transports()
    }

    fn learn_relays(&self, addresses: &[String]) {
        self.inner.learn_relays(addresses)
    }

    fn destination(&self, peer: &str) -> Option<crate::transport::Destination> {
        self.inner.destination(peer)
    }

    fn connection(&self, target: &str) -> Option<crate::transport::Connection> {
        self.inner.connection(target)
    }

    fn stats(&self) -> crate::transport::stats::TransportStats {
        self.inner.stats()
    }

    fn answers_pings(&self) -> bool {
        self.inner.answers_pings()
    }

    fn rtt(&self, peer: &str) -> Option<std::time::Duration> {
        self.inner.rtt(peer)
    }

    fn clock_offset(&self, peer: &str) -> Option<crate::transport::paths::ClockOffset> {
        self.inner.clock_offset(peer)
    }

    async fn resolve_seed(&self, seed: &str) -> Result<String> {
        self.inner.resolve_seed(seed).await
    }

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
        self.inner.encode(peer, message)
    }

    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
        let result = self.inner.send_encoded(message).await;
        if result.is_ok() {
            self.traffic
                .sent(&message.peer, message.kind, message.payload_len);
        } else {
            self.traffic
                .refused(&message.peer, message.kind, message.payload_len);
        }
        if let Some(parsed) = &message.message {
            let way = if result.is_ok() { "sent" } else { "refused" };
            self.traffic
                .log
                .note(way, &message.peer, parsed, message.payload_len);
        }
        result
    }

    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        let traffic = Arc::clone(&self.traffic);
        self.inner
            .start_consumer(Box::new(move |sender, message, payload_len| {
                traffic.received(&sender, message.kind(), payload_len);
                traffic.log.note("received", &sender, &message, payload_len);
                handler(sender, message, payload_len)
            }))
            .await
    }

    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
        self.inner.wait_for_exit().await
    }
}

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

    #[test]
    fn totals_survive_retain() {
        let traffic = Traffic::default();
        traffic.sent("a", MessageKind::Status, 100);
        traffic.sent("b", MessageKind::Sync, 1000);
        traffic.received("b", MessageKind::StatusReply, 90);
        traffic.refused("a", MessageKind::Sync, 500);

        let b = traffic.peer("b").unwrap();
        assert_eq!(b.sent.sync.messages, 1);
        assert_eq!(b.received.status_reply.bytes, 90);

        traffic.retain(|peer| peer == "b");
        assert!(traffic.peer("a").is_none());

        let total = traffic.total();
        assert_eq!(total.sent.status.bytes, 100);
        assert_eq!(total.sent.sync.bytes, 1000);
        assert_eq!(total.refused.sync.bytes, 500);
    }
}