// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Tracking of some useful statistics.

use std::{
    cell::RefCell,
    fmt::{self, Debug},
    ops::{Deref, DerefMut},
    rc::Rc,
    time::Duration,
};

use enum_map::EnumMap;
use neqo_common::{Dscp, Ecn, qdebug};
use strum::IntoEnumIterator as _;

use crate::{cc::CongestionEvent, ecn, packet};

#[derive(Default, Clone, PartialEq, Eq)]
pub struct FrameStats {
    pub ack: usize,
    pub largest_acknowledged: packet::Number,

    pub crypto: usize,
    pub stream: usize,
    pub reset_stream: usize,
    pub stop_sending: usize,

    pub ping: usize,
    pub padding: usize,

    pub max_streams: usize,
    pub streams_blocked: usize,
    pub max_data: usize,
    pub data_blocked: usize,
    pub max_stream_data: usize,
    pub stream_data_blocked: usize,

    pub new_connection_id: usize,
    pub retire_connection_id: usize,

    pub path_challenge: usize,
    pub path_response: usize,

    pub connection_close: usize,
    pub handshake_done: usize,
    pub new_token: usize,

    pub ack_frequency: usize,
    pub datagram: usize,
}

impl Debug for FrameStats {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(
            f,
            "    crypto {} done {} token {} close {}",
            self.crypto, self.handshake_done, self.new_token, self.connection_close,
        )?;
        writeln!(
            f,
            "    ack {} (max {}) ping {} padding {}",
            self.ack, self.largest_acknowledged, self.ping, self.padding
        )?;
        writeln!(
            f,
            "    stream {} reset {} stop {}",
            self.stream, self.reset_stream, self.stop_sending,
        )?;
        writeln!(
            f,
            "    max: stream {} data {} stream_data {}",
            self.max_streams, self.max_data, self.max_stream_data,
        )?;
        writeln!(
            f,
            "    blocked: stream {} data {} stream_data {}",
            self.streams_blocked, self.data_blocked, self.stream_data_blocked,
        )?;
        writeln!(f, "    datagram {}", self.datagram)?;
        writeln!(
            f,
            "    ncid {} rcid {} pchallenge {} presponse {}",
            self.new_connection_id,
            self.retire_connection_id,
            self.path_challenge,
            self.path_response,
        )?;
        writeln!(f, "    ack_frequency {}", self.ack_frequency)
    }
}

#[cfg(test)]
impl FrameStats {
    pub const fn all(&self) -> usize {
        self.ack
            + self.crypto
            + self.stream
            + self.reset_stream
            + self.stop_sending
            + self.ping
            + self.padding
            + self.max_streams
            + self.streams_blocked
            + self.max_data
            + self.data_blocked
            + self.max_stream_data
            + self.stream_data_blocked
            + self.new_connection_id
            + self.retire_connection_id
            + self.path_challenge
            + self.path_response
            + self.connection_close
            + self.handshake_done
            + self.new_token
            + self.ack_frequency
            + self.datagram
    }
}

/// Datagram stats
#[derive(Default, Clone, PartialEq, Eq)]
pub struct DatagramStats {
    /// The number of datagrams declared lost.
    pub lost: usize,
    /// The number of datagrams dropped due to being too large.
    pub dropped_too_big: usize,
    /// The number of datagrams dropped due to reaching the limit of the
    /// outgoing queue.
    pub dropped_queue_full: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SlowStartExitReason {
    /// Exited due to a congestion event (loss or ECN).
    CongestionEvent,
    /// Exited due to a heuristic algorithm (e.g., HyStart++).
    Heuristic,
}

/// Congestion Control stats
#[derive(Default, Clone, PartialEq)]
pub struct CongestionControlStats {
    /// Total number of congestion events caused by packet loss, total number of
    /// congestion events caused by ECN-CE marked packets, and number of
    /// spurious congestion events, where congestion was incorrectly inferred
    /// due to packets initially considered lost but subsequently acknowledged.
    /// The latter indicates instances where the congestion control algorithm
    /// overreacted to perceived losses.
    pub congestion_events: EnumMap<CongestionEvent, usize>,
    /// The congestion window size (in bytes) when we exited slow start.
    /// None if we haven't exited slow start or if we re-entered after spurious congestion.
    /// When exiting via congestion event, this is the cwnd AFTER the reduction.
    pub slow_start_exit_cwnd: Option<usize>,
    /// The reason slow start was exited. None if we haven't exited slow start or if we re-entered
    /// after spurious congestion.
    pub slow_start_exit_reason: Option<SlowStartExitReason>,
    /// Number of times HyStart++ entered CSS (Conservative Slow Start). Only meaningful when
    /// HyStart++ is enabled. Higher values indicate that HyStart++ had many spurious CSS
    /// entries, spending more time throttling slow start growth.
    pub hystart_css_entries: usize,
    /// Number of CSS (Conservative Slow Start) rounds completed. Only meaningful when HyStart++ is
    /// enabled. Higher values indicate the heuristic spent more time throttling slow start growth.
    pub hystart_css_rounds_finished: usize,
    /// Cubic's `w_max`: the congestion window (in bytes) just before the most recent
    /// congestion reduction (with fast convergence applied). `None` if no congestion event has
    /// occurred or Cubic is not in use. Recorded as a stat to approximate a connection's ideal
    /// congestion window in metrics.
    pub w_max: Option<f64>,
    /// The current congestion window size (in bytes). Updated throughout the connection
    /// lifetime.
    pub cwnd: Option<usize>,
}
/// ECN counts by QUIC [`packet::Type`].
#[derive(Default, Clone, PartialEq, Eq)]
pub struct EcnCount(EnumMap<packet::Type, ecn::Count>);

impl Debug for EcnCount {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (pt, count) in self.0 {
            // Don't show all-zero rows.
            if count.is_empty() {
                continue;
            }
            writeln!(f, "      {pt:?} {count:?}")?;
        }
        Ok(())
    }
}

impl Deref for EcnCount {
    type Target = EnumMap<packet::Type, ecn::Count>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for EcnCount {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// Packet types and numbers of the first ECN mark transition between two marks.
#[derive(Default, Clone, PartialEq, Eq)]
pub struct EcnTransitions(EnumMap<Ecn, EnumMap<Ecn, Option<(packet::Type, packet::Number)>>>);

impl Deref for EcnTransitions {
    type Target = EnumMap<Ecn, EnumMap<Ecn, Option<(packet::Type, packet::Number)>>>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for EcnTransitions {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl Debug for EcnTransitions {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for from in Ecn::iter() {
            // Don't show all-None rows.
            if self.0[from].iter().all(|(_, v)| v.is_none()) {
                continue;
            }
            write!(f, "      First {from:?} ")?;
            for to in Ecn::iter() {
                // Don't show transitions that were not recorded.
                if let Some(pkt) = self.0[from][to] {
                    write!(f, "to {to:?} {pkt:?} ")?;
                }
            }
            writeln!(f)?;
        }
        Ok(())
    }
}

/// Received packet counts by DSCP value.
#[derive(Default, Clone, PartialEq, Eq)]
pub struct DscpCount(EnumMap<Dscp, usize>);

impl Debug for DscpCount {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (dscp, count) in self.0 {
            // Don't show zero counts.
            if count == 0 {
                continue;
            }
            write!(f, "{dscp:?}: {count} ")?;
        }
        Ok(())
    }
}

impl Deref for DscpCount {
    type Target = EnumMap<Dscp, usize>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for DscpCount {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// Connection statistics
#[derive(Default, Clone, PartialEq)]
pub struct Stats {
    pub info: String,

    /// Total packets received, including all the bad ones.
    pub packets_rx: usize,
    /// Duplicate packets received.
    pub dups_rx: usize,
    /// Dropped packets or dropped garbage.
    pub dropped_rx: usize,
    /// The number of packet that were saved for later processing.
    pub saved_datagrams: usize,

    /// Total packets sent.
    pub packets_tx: usize,
    /// Total number of packets that are declared lost.
    pub lost: usize,
    /// Late acknowledgments, for packets that were declared lost already.
    pub late_ack: usize,
    /// Acknowledgments for packets that contained data that was marked
    /// for retransmission when the PTO timer popped.
    pub pto_ack: usize,
    /// Number of times we had to drop an unacknowledged ACK range.
    pub unacked_range_dropped: usize,
    /// Number of PMTUD probes sent.
    pub pmtud_tx: usize,
    /// Number of PMTUD probes ACK'ed.
    pub pmtud_ack: usize,
    /// Number of PMTUD probes lost.
    pub pmtud_lost: usize,
    /// MTU of the local interface used for the most recent path.
    pub pmtud_iface_mtu: usize,
    /// The peer's `max_udp_payload_size` transport parameter.
    pub pmtud_peer_max_udp_payload: Option<usize>,
    /// Probed PMTU of the current path.
    pub pmtud_pmtu: usize,

    /// Whether the connection was resumed successfully.
    pub resumed: bool,

    /// The current, estimated round-trip time on the primary path.
    pub rtt: Duration,
    /// The current, estimated round-trip time variation on the primary path.
    pub rttvar: Duration,
    /// Whether the first RTT sample was guessed from a discarded packet.
    pub rtt_init_guess: bool,

    /// Count PTOs. Single PTOs, 2 PTOs in a row, 3 PTOs in row, etc. are counted
    /// separately.
    pub pto_counts: [usize; Self::MAX_PTO_COUNTS],

    /// Count frames received.
    pub frame_rx: FrameStats,
    /// Count frames sent.
    pub frame_tx: FrameStats,

    /// The number of incoming datagrams dropped due to reaching the limit
    /// of the incoming queue.
    pub incoming_datagram_dropped: usize,

    pub datagram_tx: DatagramStats,

    pub cc: CongestionControlStats,

    /// ECN path validation count, indexed by validation outcome.
    pub ecn_path_validation: ecn::ValidationCount,
    /// ECN counts for outgoing UDP datagrams, recorded locally. For coalesced packets,
    /// counts increase for all packet types in the coalesced datagram.
    pub ecn_tx: EcnCount,
    /// ECN counts for outgoing UDP datagrams, returned by remote through QUIC ACKs.
    ///
    /// Note: Given that QUIC ACKs only carry [`Ect0`], [`Ect1`] and [`Ce`], but
    /// never [`NotEct`], the [`NotEct`] value will always be 0.
    ///
    /// See also <https://www.rfc-editor.org/rfc/rfc9000.html#section-19.3.2>.
    ///
    /// [`Ect0`]: neqo_common::tos::Ecn::Ect0
    /// [`Ect1`]: neqo_common::tos::Ecn::Ect1
    /// [`Ce`]: neqo_common::tos::Ecn::Ce
    /// [`NotEct`]: neqo_common::tos::Ecn::NotEct
    pub ecn_tx_acked: EcnCount,
    /// ECN counts for incoming UDP datagrams, read from IP TOS header. For coalesced packets,
    /// counts increase for all packet types in the coalesced datagram.
    pub ecn_rx: EcnCount,
    /// Packet numbers of the first observed (received) ECN mark transition between two marks.
    pub ecn_last_mark: Option<Ecn>,
    pub ecn_rx_transition: EcnTransitions,

    /// Counters for DSCP values received.
    pub dscp_rx: DscpCount,
}

impl Stats {
    pub const MAX_PTO_COUNTS: usize = 16;

    pub fn init(&mut self, info: String) {
        self.info = info;
    }

    pub fn pkt_dropped<A: AsRef<str>>(&mut self, reason: A) {
        self.dropped_rx += 1;
        qdebug!(
            "[{}] Dropped received packet: {}; Total: {}",
            self.info,
            reason.as_ref(),
            self.dropped_rx
        );
    }

    /// # Panics
    ///
    /// When preconditions are violated.
    pub fn add_pto_count(&mut self, count: usize) {
        debug_assert!(count > 0);
        if count >= Self::MAX_PTO_COUNTS {
            // We can't move this count any further, so stop.
            return;
        }
        self.pto_counts[count - 1] += 1;
        if count > 1 {
            debug_assert!(self.pto_counts[count - 2] > 0);
            self.pto_counts[count - 2] -= 1;
        }
    }
}

impl Debug for Stats {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "stats for {}", self.info)?;
        writeln!(
            f,
            "  rx: {} drop {} dup {} saved {}",
            self.packets_rx, self.dropped_rx, self.dups_rx, self.saved_datagrams
        )?;
        writeln!(
            f,
            "  tx: {} lost {} lateack {} ptoack {} unackdrop {}",
            self.packets_tx, self.lost, self.late_ack, self.pto_ack, self.unacked_range_dropped
        )?;
        writeln!(f, "  cc:")?;
        writeln!(
            f,
            "    ce_loss {} ce_ecn {} ce_spurious {}",
            self.cc.congestion_events[CongestionEvent::Loss],
            self.cc.congestion_events[CongestionEvent::Ecn],
            self.cc.congestion_events[CongestionEvent::Spurious],
        )?;
        writeln!(
            f,
            "    final_cwnd {:?} ss_exit_cwnd {:?} ss_exit_reason {:?}",
            self.cc.cwnd, self.cc.slow_start_exit_cwnd, self.cc.slow_start_exit_reason
        )?;
        writeln!(
            f,
            "  pmtud: {} sent {} acked {} lost {} iface_mtu {:?} peer_max_udp_payload {} pmtu",
            self.pmtud_tx,
            self.pmtud_ack,
            self.pmtud_lost,
            self.pmtud_iface_mtu,
            self.pmtud_peer_max_udp_payload,
            self.pmtud_pmtu
        )?;
        writeln!(f, "  resumed: {}", self.resumed)?;
        writeln!(f, "  frames rx:")?;
        self.frame_rx.fmt(f)?;
        writeln!(f, "  frames tx:")?;
        self.frame_tx.fmt(f)?;
        writeln!(f, "  ecn:\n    tx:")?;
        self.ecn_tx.fmt(f)?;
        writeln!(f, "    acked:")?;
        self.ecn_tx_acked.fmt(f)?;
        writeln!(f, "    rx:")?;
        self.ecn_rx.fmt(f)?;
        writeln!(
            f,
            "    path validation outcomes: {:?}",
            self.ecn_path_validation
        )?;
        writeln!(f, "    mark transitions:")?;
        self.ecn_rx_transition.fmt(f)?;
        writeln!(f, "  dscp: {:?}", self.dscp_rx)
    }
}

#[derive(Default, Clone)]
pub struct StatsCell {
    stats: Rc<RefCell<Stats>>,
}

impl Deref for StatsCell {
    type Target = RefCell<Stats>;
    fn deref(&self) -> &Self::Target {
        &self.stats
    }
}

impl Debug for StatsCell {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.stats.borrow().fmt(f)
    }
}

#[test]
fn debug() {
    let stats = Stats::default();
    assert_eq!(
        format!("{stats:?}"),
        "stats for\u{0020}
  rx: 0 drop 0 dup 0 saved 0
  tx: 0 lost 0 lateack 0 ptoack 0 unackdrop 0
  cc:
    ce_loss 0 ce_ecn 0 ce_spurious 0
    final_cwnd None ss_exit_cwnd None ss_exit_reason None
  pmtud: 0 sent 0 acked 0 lost 0 iface_mtu None peer_max_udp_payload 0 pmtu
  resumed: false
  frames rx:
    crypto 0 done 0 token 0 close 0
    ack 0 (max 0) ping 0 padding 0
    stream 0 reset 0 stop 0
    max: stream 0 data 0 stream_data 0
    blocked: stream 0 data 0 stream_data 0
    datagram 0
    ncid 0 rcid 0 pchallenge 0 presponse 0
    ack_frequency 0
  frames tx:
    crypto 0 done 0 token 0 close 0
    ack 0 (max 0) ping 0 padding 0
    stream 0 reset 0 stop 0
    max: stream 0 data 0 stream_data 0
    blocked: stream 0 data 0 stream_data 0
    datagram 0
    ncid 0 rcid 0 pchallenge 0 presponse 0
    ack_frequency 0
  ecn:
    tx:
    acked:
    rx:
    path validation outcomes: ValidationCount({Capable: 0, NotCapable(BlackHole): 0, NotCapable(Bleaching): 0, NotCapable(ReceivedUnsentECT1): 0})
    mark transitions:
  dscp: \n"
    );
}