//! Allows for the logging of user-defined metrics and events.
//! The simulator also records internal stats and events which may be disabled with the `disable_internal_stats` and `disable_internal_events` package features.
use crate::node::NodeID;
use linearize::{Linearize, StaticMap};
use serde::Serialize;

/// Key for stats that the simulator tracks internally.
#[derive(Linearize, Serialize)]
pub enum InternalStatKey {
    /// Number of packets transmitted across a tick.
    PacketTransmits,
    /// Number of packets received across all nodes in a tick.
    /// Can be more than transmitted as each packet can be received by multiple nodes.
    PacketReceives,
}

/// Events generated by the library itself.
#[derive(Clone, Debug, Serialize)]
pub enum InternalEvent {
    /// A node transmitted a packet
    PacketTransmit(NodeID),
    /// A src sent a packet to a dst
    PacketLink((NodeID, NodeID)),
}

/// Used to record arbitrary stats for a given timestep
/// Note you cannot retrieve stats during a tick, as stats are thread-localised
pub struct TimestepStats<T: Linearize, E: Clone> {
    internal: StaticMap<InternalStatKey, isize>,
    pub(crate) internal_events: Vec<InternalEvent>,
    user: StaticMap<T, isize>,
    user_events: Vec<E>,
}

impl<T: Linearize, E: Clone> TimestepStats<T, E> {
    pub(crate) fn new() -> Self {
        Self {
            internal: StaticMap::default(),
            user: StaticMap::default(),
            internal_events: Vec::new(),
            user_events: Vec::new(),
        }
    }

    /// Increment a user-defined stat.
    pub fn inc(&mut self, key: T, by: isize) {
        self.user[key] += by;
    }

    /// Decrement a user-defined stat.
    pub fn dec(&mut self, key: T, by: isize) {
        self.user[key] -= by;
    }

    #[cfg(not(feature = "disable_internal_stats"))]
    pub(crate) fn inc_internal(&mut self, key: InternalStatKey, by: isize) {
        self.internal[key] += by;
    }

    #[cfg(not(feature = "disable_internal_stats"))]
    pub(crate) fn dec_internal(&mut self, key: InternalStatKey, by: isize) {
        self.internal[key] -= by;
    }

    #[cfg(not(feature = "disable_internal_events"))]
    pub(crate) fn add_internal_event(&mut self, event: InternalEvent) {
        self.internal_events.push(event);
    }

    pub(crate) fn add_user_event(&mut self, event: E) {
        self.user_events.push(event);
    }

    /// Merges another [TimestepStats] into this one, leaving the other one empty.
    pub(crate) fn consume(&mut self, other: Self) {
        #[cfg(not(feature = "disable_internal_stats"))]
        for (key, val) in &other.internal {
            self.internal[key] += val;
        }

        for (key, val) in &other.user {
            self.user[key] += val;
        }

        #[cfg(not(feature = "disable_internal_events"))]
        self.internal_events
            .extend(other.internal_events.into_iter());

        self.user_events.extend(other.user_events.into_iter())
    }

    /// Consume and return events buffer, buffer will be left empty afterwards.
    pub fn events(&mut self) -> Vec<E> {
        std::mem::take(&mut self.user_events)
    }

    #[cfg(not(feature = "disable_internal_events"))]
    /// Consume and return internal event buffer, will be left empty afterwards.
    pub fn internal_events(&mut self) -> Vec<InternalEvent> {
        std::mem::take(&mut self.internal_events)
    }

    /// Consume and return stats, will be left empty afterwards.
    pub fn stats(&mut self) -> StaticMap<T, isize> {
        std::mem::take(&mut self.user)
    }

    #[cfg(not(feature = "disable_internal_stats"))]
    /// Consume and return internal stats, will be left empty afterwards.
    pub fn internal_stats(&mut self) -> StaticMap<InternalStatKey, isize> {
        std::mem::take(&mut self.internal)
    }
}

impl<T: Linearize, E: Clone> Default for TimestepStats<T, E> {
    fn default() -> Self {
        TimestepStats::new()
    }
}
