// BgpSim: BGP Network Simulator written in Rust
// Copyright 2022-2024 Tibor Schneider <sctibor@ethz.ch>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Module to record actions on the network and extract a trace of how the forwarding state changes
//! over time.

use serde::{Deserialize, Serialize};

use crate::{
    event::EventQueue,
    forwarding_state::ForwardingState,
    interactive::InteractiveNetwork,
    network::Network,
    types::{NetworkError, RouterId, SinglePrefix, StepUpdate},
};

/// Extension trait that allows you to record events on the network. This is only available for
/// [`SinglePrefix`].
pub trait RecordNetwork<Q> {
    /// Simulate a single action on the network and extract the record of the convergence
    /// process. This function will record the forwarding state of all prefixes in the network. In
    /// addition, if applicable, the times when the forwarding state changes will be recorded.
    ///
    /// This function will return either the trace and the initial forwarding state, or the network
    /// error generated by `f`.
    fn record<F>(&mut self, f: F) -> Result<ConvergenceRecording, NetworkError>
    where
        F: FnOnce(&mut Network<SinglePrefix, Q>) -> Result<(), NetworkError>;

    /// Simulate the network (by running all enqueued events). During this, construct a
    /// `ConvergenceRecording` by recording every single change in the forwarding behavior of the
    /// network. This function takes `initial_fw_state` (before any event happened), and the
    /// `initial_trace`, that contains the difference from the initial forarding state to the state
    /// that is currently present. The `initial_time` is used to normalize the provided convergence
    /// times, e.g., to make every ConvergenceRecording zero-normalized.
    ///
    /// Consider the implementation of [`RecordNetwork::record`] for `Network` to see an example of
    /// how to use this function.
    fn record_prepared(
        &mut self,
        initial_fw_state: ForwardingState<SinglePrefix>,
        initial_trace: ConvergenceTrace,
        initial_time: Option<f64>,
    ) -> Result<ConvergenceRecording, NetworkError>;
}

impl<Q> RecordNetwork<Q> for Network<SinglePrefix, Q>
where
    Q: EventQueue<SinglePrefix>,
{
    fn record<F>(&mut self, f: F) -> Result<ConvergenceRecording, NetworkError>
    where
        F: FnOnce(&mut Network<SinglePrefix, Q>) -> Result<(), NetworkError>,
    {
        // set the queue to skip mode.
        let original_skip_value = self.skip_queue;
        self.skip_queue = true;

        // get the forwarding state before
        let fw_state_before = self.get_forwarding_state();

        let initial_time = self.queue().get_time();

        // execute the function
        f(self)?;

        // get the forwarding state difference and start generating the trace
        let fw_state_after = self.get_forwarding_state();
        let diff = fw_state_before.diff(&fw_state_after);

        let trace = vec![(diff, initial_time.map(|_| 0.0).into())];

        // now, generate the recording using the intrinsic function
        let record = self.record_prepared(fw_state_before, trace, initial_time)?;

        // reset the queue skip property
        self.skip_queue = original_skip_value;

        Ok(record)
    }

    fn record_prepared(
        &mut self,
        initial_fw_state: ForwardingState<SinglePrefix>,
        mut trace: ConvergenceTrace,
        initial_time: Option<f64>,
    ) -> Result<ConvergenceRecording, NetworkError> {
        let t = initial_time.unwrap_or_default();
        while let Some((step, event)) = self.simulate_step()? {
            match step {
                StepUpdate::Unchanged => {}
                StepUpdate::Single(delta) => {
                    let time = self.queue().get_time().map(|x| x - t);
                    trace.push((vec![(event.router(), delta.old, delta.new)], time.into()));
                }
                StepUpdate::Multiple => {
                    log::warn!("Ignoring OSPF event that might changed multiple FIB entries!")
                }
            }
        }
        Ok(ConvergenceRecording::new(initial_fw_state, trace))
    }
}

/// Record of an entire convergence process that captures the history of each change in the
/// forwarding state of the network. This is a record that allows the state to be moved forwards or
/// backwards in time. The recording can be changed on a per-prefix level.
#[derive(Debug, Clone, PartialEq)]
pub struct ConvergenceRecording {
    state: ForwardingState<SinglePrefix>,
    trace: ConvergenceTrace,
    pointer: usize,
}

impl ConvergenceRecording {
    /// Create a Recording from a trace and an initial forwarding state
    pub fn new(initial_fw_state: ForwardingState<SinglePrefix>, trace: ConvergenceTrace) -> Self {
        ConvergenceRecording {
            state: initial_fw_state,
            trace,
            pointer: 0,
        }
    }

    /// Get a reference to the current forwarding state
    pub fn state(&mut self) -> &mut ForwardingState<SinglePrefix> {
        &mut self.state
    }

    /// Get a reference of the convergence trace.
    pub fn trace(&self) -> &ConvergenceTrace {
        &self.trace
    }

    /// Transform the recording into a trace.
    pub fn as_trace(self) -> ConvergenceTrace {
        self.trace
    }

    /// Perform a single step for an individual prefix. If the forwarding state is already in the
    /// final state for the specifiied prefix, then this function will return `None`. Otherwise, it
    /// will return a slice containing all deltas that were applied during this function call, the
    /// network's convergence time when the change took effect (if applicable), and a mutable
    /// reference to the new `ForwardingState`.
    pub fn step(
        &mut self,
    ) -> Option<(&[FwDelta], Option<f64>, &mut ForwardingState<SinglePrefix>)> {
        if self.pointer >= self.trace.len() {
            // already out of bounds
            return None;
        }
        // apply the delta at the current position
        let (deltas, time) = self.trace.get(self.pointer)?;
        for (router, _, new_nh) in deltas {
            self.state.update(*router, SinglePrefix, new_nh.clone());
        }

        // increment the pointer
        self.pointer += 1;

        // return the applied deltas
        Some((deltas, time.into_inner(), &mut self.state))
    }

    /// Undo a single step for an individual prefix. If the forwarding state is already in the
    /// initial state for the specifiied prefix, then this function will return `None`. Otherwise, it
    /// will return a slice containing all deltas that were applied *in reverse direction* during
    /// this function call, the network's convergence time when the change took effect (if
    /// applicable), and a mutable reference to the new `ForwardingState`.
    pub fn back(
        &mut self,
    ) -> Option<(&[FwDelta], Option<f64>, &mut ForwardingState<SinglePrefix>)> {
        if self.pointer == 0 {
            // already out of bounds
            return None;
        }
        // decrement the pointer
        self.pointer -= 1;
        // apply the delta at the current position
        let (deltas, time) = self.trace.get(self.pointer)?;
        for (router, old_nh, _) in deltas {
            self.state.update(*router, SinglePrefix, old_nh.clone());
        }

        // return the applied deltas
        Some((deltas, time.into_inner(), &mut self.state))
    }

    /// Get the position of the recording for the given prefix.
    pub fn pos(&self) -> usize {
        self.pointer
    }

    /// Get the length of the recording for the given prefix
    #[allow(clippy::len_without_is_empty)]
    pub fn len(&self) -> usize {
        self.trace.len()
    }

    /// Check if the recording is empty for the given prefix
    pub fn is_empty(&self) -> bool {
        self.trace.is_empty()
    }

    /// Reverts to and releases the initial `ForwardingState` held by `self`, while consuming
    /// `self`. This function will also drop any timing information.
    pub fn into_initial_fw_state(self) -> ForwardingState<SinglePrefix> {
        let mut state = self.state;
        let trace = self.trace;
        let mut ptr = self.pointer;

        while ptr > 0 {
            ptr -= 1;
            let (deltas, _) = trace.get(ptr).unwrap();
            for (router, old_nh, _) in deltas {
                state.update(*router, SinglePrefix, old_nh.clone());
            }
        }
        state
    }
}

/// This structure captures the essence of a trace, that is, the entire evolution of the forwarding
/// state during the convergence process. It does not capture the initial or the final state, but it
/// should be efficient to compare entries.
///
/// This structure also stores the time when the forwarding delta occurs. The time is
/// zero-normalized, which means that the time zero always represents the time when the convergence
/// trace recording has started. The time is warpped in an `AlwaysEq` such that comparisons will
/// ignore the time.
pub type ConvergenceTrace = Vec<(Vec<FwDelta>, AlwaysEq<Option<f64>>)>;

/// Structure that wraps the inner type and always returns `true` or `Ordering::Equal` when compared.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[repr(transparent)]
pub struct AlwaysEq<T>(T);

impl<T> std::hash::Hash for AlwaysEq<T> {
    fn hash<H: std::hash::Hasher>(&self, _state: &mut H) {}
}

impl<T> Ord for AlwaysEq<T> {
    fn cmp(&self, _other: &Self) -> std::cmp::Ordering {
        std::cmp::Ordering::Equal
    }
}

impl<T> PartialOrd for AlwaysEq<T> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<T> Eq for AlwaysEq<T> {}

impl<T> PartialEq for AlwaysEq<T> {
    fn eq(&self, _other: &Self) -> bool {
        true
    }
}

impl<T> From<T> for AlwaysEq<T> {
    fn from(value: T) -> Self {
        AlwaysEq(value)
    }
}

impl<T> From<AlwaysEq<Option<T>>> for Option<T> {
    fn from(value: AlwaysEq<Option<T>>) -> Self {
        value.0
    }
}

impl<T> AsRef<T> for AlwaysEq<T> {
    fn as_ref(&self) -> &T {
        &self.0
    }
}

impl<T: std::fmt::Display> std::fmt::Display for AlwaysEq<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

impl<T> AlwaysEq<T> {
    /// Consume `self` and return the inner `T`.
    pub fn into_inner(self) -> T {
        self.0
    }
}

/// Forwarding state delta.
pub type FwDelta = (RouterId, Vec<RouterId>, Vec<RouterId>);
