//! 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 the Mycelium message carries, without its framing,
//! 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 still counts as sent (the firewall plays a
//! lossy network) but not as received.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};

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

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

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

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

#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct ByKind {
    pub hello: Counter,
    pub hello_reply: Counter,
    pub sync: Counter,
    pub ping: Counter,
    pub pong: 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::Ping => &mut self.ping,
            MessageKind::Pong => &mut self.pong,
        }
    }
}

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

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

#[derive(Debug, Default)]
pub struct Traffic {
    table: Mutex<Table>,
}

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

    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),
        }
    }

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

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

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

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

    pub 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 fn retain(&self, mut keep: impl FnMut(&str) -> bool) {
        self.lock().peers.retain(|peer, _| keep(peer));
    }
}

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

impl CountingTransport {
    pub 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()
    }

    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);
        }
        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);
                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::Ping, 100);
        traffic.sent("b", MessageKind::Sync, 1000);
        traffic.received("b", MessageKind::Pong, 90);
        traffic.refused("a", MessageKind::Sync, 500);

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

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

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