// 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.

//! This module contains an extension trait that allows you to interact with the simulator on a
//! per-message level.

use crossbeam::channel::{bounded, Receiver, Sender};
use std::collections::BTreeMap;

use log::debug;

use crate::{
    event::{ConcurrentEventQueue, Event, EventQueue},
    formatter::NetworkFormatter,
    network::Network,
    ospf::{global::GlobalOspf, OspfImpl, OspfProcess},
    router::Router,
    types::{DeviceError, NetworkError, Prefix, RouterId, StepUpdate},
};

/// Trait that allows you to interact with the simulator on a per message level. It exposes an
/// interface to simulate a single event, inspect the queue of the network, and even reorder events.
pub trait InteractiveNetwork<P, Q, Ospf>
where
    P: Prefix,
    Q: EventQueue<P>,
    Ospf: OspfImpl,
{
    /// Setup the network to automatically simulate each change of the network. This is the default
    /// behavior. Disable auto-simulation by using [`InteractiveNetwork::manual_simulation`].
    fn auto_simulation(&mut self);

    /// Setup the network to not to automatically simulate each change of the network. Upon any
    /// change of the network (configuration change, external update of any routing input, or a link
    /// failure), the event queue will be filled with the initial message(s), but it will not
    /// execute them. Enable auto-simulation by using [`InteractiveNetwork::auto_simulation`]. Use
    /// either [`Network::simulate`] to run the entire queue after updating the messages, or use
    /// [`InteractiveNetwork::simulate_step`] to execute a single event on the queue.
    fn manual_simulation(&mut self);

    /// Returns `true` if auto-simulation is enabled.
    fn auto_simulation_enabled(&self) -> bool;

    /// Calls the function `f` with argument to a mutable network. During this call, the network
    /// will have automatic simulation disabled. It will be re-enabled once the function exits.
    ///
    /// Note, that this function takes ownership of `self` and returns it afterwards. This is to
    /// prohibit you to call `with_manual_simulation` multiple times.
    fn with_manual_simulation<F>(self, f: F) -> Self
    where
        F: FnOnce(&mut Self);

    /// Simulate the network behavior, given the current event queue. This function will execute all
    /// events (that may trigger new events), until either the event queue is empt (i.e., the
    /// network has converged), or until the maximum allowed events have been processed (which can
    /// be set by `self.set_msg_limit`).
    fn simulate(&mut self) -> Result<(), NetworkError>;

    /// Similarly to the [`Network::simulate`] function, this function will execute all events in the
    /// queue. This function provides a hook in the form of a closure that is called before and after
    /// each event is processed.
    ///
    /// The closure always provides *immutable* access to both the network and the currently processed event.
    /// In order to determine whether it was called before or after the event was popped from the queue and
    /// processed we need to look at the value of the third argument.
    ///
    /// If the closure is called *before* the event is processed:
    ///  - The third argument will be `None`
    ///  - The event is still present in the network's queue
    ///
    /// If the closure is called *after* the event is processed:
    ///  - The third argument will contain a `Some((StepUpdate, Vec<Event>))` tuple, where the second element
    ///    contains a list of events that were generated by the event that was just processed.
    ///  - These events have not been enqueued yet
    fn simulate_hooked(
        &mut self,
        f: impl FnMut(
            &Self,
            &Event<P, Q::Priority>,
            Option<&(StepUpdate<P>, Vec<Event<P, Q::Priority>>)>,
        ),
    ) -> Result<(), NetworkError>;

    /// Trigger the timeout event on any router. The router is picked randomly if the feature `rand`
    /// is enabled. The function returns the router on which the timeout was triggered, or `None` if
    /// no router is waiting for a timeout event.
    ///
    /// After calling this function, the queue might contain new events. Run `net.simulate()` to
    /// execute them.
    ///
    /// Timeout might cause OSPF events to be generated at internal routers.
    fn trigger_timeout(&mut self) -> Result<Option<RouterId>, NetworkError>;

    /// Trigger the timeout event on `router`. If the router is not waiting for a timeout event, the
    /// function returns `Ok(false)`. Otherwise, the timeout is triggered, and `Ok(true)` is
    /// returned.
    ///
    /// After calling this function, the queue might contain new events. Run `net.simulate()` to
    /// execute them.
    ///
    /// Timeout might cause OSPF events to be generated at internal routers.
    fn trigger_timeout_at(&mut self, router: RouterId) -> Result<bool, NetworkError>;

    /// Simulate the next event on the queue. In comparison to [`Network::simulate`], this function
    /// will not execute any subsequent event. This function returns the change in forwarding
    /// behavior caused by this step, as well as the event that was processed. If this function
    /// returns `Ok(None)`, then no event was enqueued.
    #[allow(clippy::type_complexity)]
    fn simulate_step(
        &mut self,
    ) -> Result<Option<(StepUpdate<P>, Event<P, Q::Priority>)>, NetworkError>;

    /// Get a reference to the queue
    fn queue(&self) -> &Q;

    /// Get a reference to the queue
    fn queue_mut(&mut self) -> &mut Q;

    /// Manually trigger the given event, returning the result of that event. No new events will be
    /// enqueued.
    ///
    /// # Safety
    /// The network will be in an inconsistent state. Make sure to deal with that properly.
    #[allow(clippy::type_complexity)]
    unsafe fn trigger_event(
        &mut self,
        event: Event<P, Q::Priority>,
    ) -> Result<(StepUpdate<P>, Vec<Event<P, Q::Priority>>), NetworkError>;

    /// Manually enqueue a specific event. The event will be executed automatically if the network
    /// is the automatic simulation state.
    ///
    /// # Safety
    /// The network will be in an inconsistent state. Make sure to deal with that properly.
    unsafe fn enqueue_event(&mut self, event: Event<P, Q::Priority>);
}

/// Trait that allows you to run simulations in parallel.
pub trait ParallelNetwork<P, Q, Ospf>
where
    P: Prefix + Send + Sync,
    Q: ConcurrentEventQueue<P>,
    Ospf: OspfImpl,
{
    /// Simulate the network behavior, given the current event queue. This function will execute all
    /// events (that may trigger new events), until either the event queue is empt (i.e., the
    /// network has converged), or until the maximum allowed events have been processed (which can
    /// be set by `self.set_msg_limit`).
    ///
    /// This function will use multiple threads to process the network concurrently.
    fn simulate_parallel(&mut self, num_threads: Option<usize>) -> Result<(), NetworkError>;
}

impl<P: Prefix, Q: EventQueue<P>, Ospf: OspfImpl> InteractiveNetwork<P, Q, Ospf>
    for Network<P, Q, Ospf>
{
    fn auto_simulation(&mut self) {
        self.skip_queue = false;
    }

    fn manual_simulation(&mut self) {
        self.skip_queue = true;
    }

    fn auto_simulation_enabled(&self) -> bool {
        !self.skip_queue
    }

    fn with_manual_simulation<F>(mut self, f: F) -> Self
    where
        F: FnOnce(&mut Self),
    {
        self.manual_simulation();
        f(&mut self);
        self.auto_simulation();
        self
    }

    fn simulate_step(
        &mut self,
    ) -> Result<Option<(StepUpdate<P>, Event<P, Q::Priority>)>, NetworkError> {
        if let Some(event) = self.queue.pop() {
            // log the job
            log::trace!("{}", event.fmt(self));
            // execute the event
            let (step_update, events) = self
                .routers
                .get_mut(&event.router())
                .ok_or(NetworkError::DeviceNotFound(event.router()))?
                .handle_event(event.clone())?;

            self.enqueue_events(events);

            Ok(Some((step_update, event)))
        } else {
            Ok(None)
        }
    }

    fn queue(&self) -> &Q {
        &self.queue
    }

    fn queue_mut(&mut self) -> &mut Q {
        &mut self.queue
    }

    fn simulate(&mut self) -> Result<(), NetworkError> {
        let mut remaining_iter = self.stop_after;
        'timeout: loop {
            while !self.queue.is_empty() {
                if let Some(rem) = remaining_iter {
                    if rem == 0 {
                        debug!("Network could not converge!");
                        return Err(NetworkError::NoConvergence);
                    }
                    remaining_iter = Some(rem - 1);
                }
                let step_update = self.simulate_step()?;
                if matches!(step_update, Some((_, Event::Ospf { .. }))) {
                    // OSPF event received! Check the BGP session state
                    self.refresh_bgp_sessions()?;
                }
            }

            // trigger the next timeout event if it exists.
            let trigger_router = self.trigger_timeout()?;

            // if no timeout was triggered, break out of the loop. We are converged!
            if trigger_router.is_none() {
                break 'timeout;
            }
        }

        // remove unreachable OSPF LSAs
        self.routers_mut()
            .for_each(|r| r.ospf.remove_unreachable_lsas());

        Ok(())
    }

    fn simulate_hooked(
        &mut self,
        mut f: impl FnMut(
            &Self,
            &Event<P, Q::Priority>,
            Option<&(StepUpdate<P>, Vec<Event<P, Q::Priority>>)>,
        ),
    ) -> Result<(), NetworkError> {
        let mut remaining_iter = self.stop_after;
        'timeout: loop {
            // While there are events in the queue
            while let Some(event) = self.queue().peek() {
                // Ensure the convergence limit is not overstepped
                if let Some(rem) = remaining_iter {
                    if rem == 0 {
                        debug!("Network could not converge!");
                        return Err(NetworkError::NoConvergence);
                    }
                    remaining_iter = Some(rem - 1);
                }

                // Straddle the trigger_event function with the pre- and post-event hooks
                f(self, event, None);
                // Safety: This is safe because we trigger the next event in the queue and
                // we still push all resulting events to the queue. The extracted events are
                // all immutable and won't be modified.
                // Since the closure's reference to the network is immutable as well, this unwrap
                // is always safe
                let event = self.queue_mut().pop().unwrap();
                let event_clone = event.clone();
                let result = unsafe { self.trigger_event(event)? };
                f(self, &event_clone, Some(&result));

                self.enqueue_events(result.1);

                if matches!(event_clone, Event::Ospf { .. }) {
                    // OSPF event received! Check the BGP session state
                    self.refresh_bgp_sessions()?;
                }
            }

            // trigger the next timeout event if it exists.
            let trigger_router = self.trigger_timeout()?;

            // if no timeout was triggered, break out of the loop. We are converged!
            if trigger_router.is_none() {
                break 'timeout;
            }
        }

        // remove unreachable OSPF LSAs
        self.routers_mut()
            .for_each(|r| r.ospf.remove_unreachable_lsas());

        Ok(())
    }

    fn trigger_timeout(&mut self) -> Result<Option<RouterId>, NetworkError> {
        #[allow(unused_mut)]
        let mut routers_waiting_for_timeout = self
            .routers()
            .filter(|r| r.is_waiting_for_timeout())
            .map(|r| r.router_id());

        #[cfg(not(feature = "rand"))]
        let router: Option<RouterId> = routers_waiting_for_timeout.next();

        #[cfg(feature = "rand")]
        let router: Option<RouterId> = {
            use rand::prelude::*;
            let mut rng = thread_rng();
            let waiting = routers_waiting_for_timeout.collect::<Vec<_>>();
            waiting.as_slice().choose(&mut rng).copied()
        };

        if let Some(router) = router {
            Ok(self.trigger_timeout_at(router)?.then_some(router))
        } else {
            Ok(None)
        }
    }

    fn trigger_timeout_at(&mut self, router: RouterId) -> Result<bool, NetworkError> {
        let r = self.get_router_mut(router)?;

        if r.is_waiting_for_timeout() {
            log::debug!("Trigger timeout on {}", r.name());
            let events = r.trigger_timeout()?;
            self.enqueue_events(events);
            Ok(true)
        } else {
            Ok(false)
        }
    }

    unsafe fn trigger_event(
        &mut self,
        event: Event<P, Q::Priority>,
    ) -> Result<(StepUpdate<P>, Vec<Event<P, Q::Priority>>), NetworkError> {
        // log the job
        log::trace!("{}", event.fmt(self));
        // execute the event
        let (step_update, events) = self
            .routers
            .get_mut(&event.router())
            .ok_or(NetworkError::DeviceNotFound(event.router()))?
            .handle_event(event.clone())?;

        Ok((step_update, events))
    }

    unsafe fn enqueue_event(&mut self, event: Event<P, Q::Priority>) {
        self.enqueue_events(vec![event]);
    }
}

impl<P, Q, Ospf> ParallelNetwork<P, Q, Ospf> for Network<P, Q, Ospf>
where
    P: Prefix,
    Q: ConcurrentEventQueue<P>,
    Q::Priority: Send + Sync + 'static,
    Ospf: OspfImpl,
{
    fn simulate_parallel(&mut self, num_threads: Option<usize>) -> Result<(), NetworkError> {
        let mut remaining_iter = self.stop_after;

        // take out all routers.
        let mut routers = BTreeMap::new();
        std::mem::swap(&mut routers, &mut self.routers);
        let mut routers = routers.into_values().collect::<Vec<_>>();

        // get the number of threads.
        let num_threads = num_threads
            .unwrap_or_else(|| {
                std::thread::available_parallelism()
                    .map(|x| x.into())
                    .unwrap_or(4)
            })
            .max(1) // at least one worker
            .min(routers.len()); // at most as many workers as there are routers
        let num_concurrent_jobs = num_threads;

        // create two channels, bounded by the number of routers in the network
        let (task_sender, task_receiver) =
            bounded::<Job<P, Q::Priority, Ospf::Process>>(routers.len() + 1);
        let (result_sender, result_receiver) =
            bounded::<Res<P, Q::Priority, Ospf::Process>>(routers.len() + 1);

        let workers = (0..num_threads)
            .map(|_| {
                let task_receiver = task_receiver.clone();
                let result_sender = result_sender.clone();
                std::thread::spawn(move || worker(task_receiver, result_sender))
            })
            .collect::<Vec<_>>();
        // drop our end of the results sender.
        std::mem::drop(result_sender);

        let num_routers = routers.len();
        let mut result = Ok(());
        let empty = BTreeMap::<RouterId, Router<P, Ospf::Process>>::new();

        // start the process
        'main: loop {
            // exit condition
            if self.queue.is_empty() && routers.len() == num_routers {
                break 'main;
            }

            // first, try to perform as many jobs as possible.
            let mut i = 0;
            while i < routers.len() {
                // Only ever enqueue as many jobs as there are threads to increase batching (plus
                // one, so one thread can immediately continue working.)
                if num_routers - routers.len() >= num_concurrent_jobs {
                    break;
                }
                // break out if queue is empty
                if self.queue.is_empty() {
                    break;
                }
                // try to trigger all events on router i
                let rid = routers[i].router_id();
                let events = self.queue.pop_events_for(rid);
                if !events.is_empty() {
                    // enqueue job
                    if let Some(x) = remaining_iter.as_mut() {
                        if *x < events.len() {
                            // no convergence!
                            result = Err(NetworkError::NoConvergence);
                            break 'main;
                        }
                        *x -= events.len();
                    }
                    let r = routers.swap_remove(i);
                    task_sender.send((r, events)).unwrap();
                }
                i += 1;
            }

            // now, receive the next task
            let (r, events, success) = result_receiver.recv().unwrap();
            routers.push(r);
            self.queue.push_many(events, &empty, &self.net);
            result = result.and(success.map_err(NetworkError::DeviceError));

            // check if we have an error
            if result.is_err() {
                break 'main;
            }
        }

        // kill all processes by dropping the task channel
        std::mem::drop(task_sender);
        workers.into_iter().for_each(|t| t.join().unwrap());

        // get the last events back to potentially reconstruct the network. Since we have already
        // dropped our part of the result sender, we can just wait here until the result channel is
        // disconnected (which will happen only after all workers have exited and all results have
        // been pulled.)
        while let Ok((r, events, success)) = result_receiver.recv() {
            routers.push(r);
            self.queue.push_many(events, &empty, &self.net);
            result = result.and(success.map_err(NetworkError::DeviceError));
        }

        // put the routers back
        self.routers = routers.into_iter().map(|r| (r.router_id(), r)).collect();

        result
    }
}

type Job<P, T, Ospf> = (Router<P, Ospf>, Vec<Event<P, T>>);
type Res<P, T, Ospf> = (Router<P, Ospf>, Vec<Event<P, T>>, Result<(), DeviceError>);

fn worker<P: Prefix, T: Default, Ospf: OspfProcess>(
    recv: Receiver<Job<P, T, Ospf>>,
    send: Sender<Res<P, T, Ospf>>,
) {
    loop {
        let Ok((mut r, events)) = recv.recv() else {
            break;
        };
        let mut result = Ok(());
        let mut out_events = Vec::new();
        for e in events {
            match r.handle_event(e) {
                Ok((_, next_events)) => out_events.extend(next_events),
                Err(e) => {
                    result = Err(e);
                    break;
                }
            };
        }

        // send the event back
        send.send((r, out_events, result)).unwrap();
    }
}

/// Builder interface to partially clone the source network while moving values from the conquered
/// network. most of the functions in this structure are `unsafe`, because the caller must guarantee
/// that the source and the conquered network share the exact same state for those values that you
/// decide to reuse.
///
/// If you do not reuse anything of the conquered network, then this function will most likely be
/// slower than simply calling `source.clone()`.
///
/// Currently, this structure only supports cloning a `Network<P, Q, GlobalOspf>`.
///
/// ```
/// # #[cfg(feature = "topology_zoo")]
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # use bgpsim::prelude::*;
/// # use bgpsim::types::SimplePrefix as P;
/// # use bgpsim::topology_zoo::TopologyZoo;
/// # use bgpsim::event::BasicEventQueue;
/// # use bgpsim::builder::*;
/// # let mut net: Network<_, _, GlobalOspf> = TopologyZoo::Abilene.build(BasicEventQueue::new(), 65500, 1);
/// # let prefix = P::from(0);
/// # net.build_external_routers(65500, 1, RandomRouters::new(3))?;
/// # net.build_ibgp_route_reflection(HighestDegreeRouters::new(2))?;
/// # net.build_ebgp_sessions()?;
/// # net.build_link_weights(20.0)?;
/// # let ads = net.build_advertisements(prefix, UniquePreference::new().internal_asn(65500), ASN(100))?;
/// # let ext = ads[0].0;
/// use bgpsim::interactive::PartialClone;
///
/// // let mut net = ...
/// let original_net = net.clone();
/// net.withdraw_route(ext, prefix)?;
/// assert_ne!(net, original_net);
/// let net = unsafe {
///     PartialClone::new(&original_net)
///         .reuse_config(true)
///         .reuse_igp_state(true)
///         .reuse_queue_params(true)
///         .conquer(net)
/// };
/// assert_eq!(net, original_net);
/// # Ok(())
/// # }
/// # #[cfg(not(feature = "topology_zoo"))]
/// # fn main() {}
/// ```
#[derive(Debug)]
pub struct PartialClone<'a, P: Prefix, Q> {
    source: &'a Network<P, Q, GlobalOspf>,
    reuse_config: bool,
    reuse_advertisements: bool,
    reuse_igp_state: bool,
    reuse_bgp_state: bool,
    reuse_queue_params: bool,
}

impl<'a, P: Prefix, Q> PartialClone<'a, P, Q> {
    /// Create a new partial clone of a network
    pub fn new(net: &'a Network<P, Q, GlobalOspf>) -> PartialClone<'a, P, Q> {
        PartialClone {
            source: net,
            reuse_igp_state: false,
            reuse_bgp_state: false,
            reuse_config: false,
            reuse_advertisements: false,
            reuse_queue_params: false,
        }
    }

    /// Reuse the entire configuration from the conquered network.
    ///
    /// # Safety
    /// The caller must ensure that the entire configuration of both the source network and the
    /// conquered network is identical.
    pub unsafe fn reuse_config(mut self, b: bool) -> Self {
        self.reuse_config = b;
        self
    }

    /// Reuse all external advertisements.
    ///
    /// # Safety
    /// The caller must ensure that the advertisements of both the source and the conquered network
    /// is identical.
    pub unsafe fn reuse_advertisements(mut self, b: bool) -> Self {
        self.reuse_advertisements = b;
        self
    }

    /// Reuse the IGP state of the network. This function requires that you also reuse the
    /// configuration!
    ///
    /// # Safety
    /// The caller must ensure that the entire IGP state of both the source and the conquered
    /// network is identical.
    pub unsafe fn reuse_igp_state(mut self, b: bool) -> Self {
        self.reuse_igp_state = b;
        self
    }

    /// Reuse the BGP state of the network. This function requires that you also reuse the
    /// configuration and the advertisements!
    ///
    /// # Safety
    /// The caller must ensure that the entire BGP state of both the source and the conquered
    /// network is identical.
    pub unsafe fn reuse_bgp_state(mut self, b: bool) -> Self {
        self.reuse_bgp_state = b;
        self
    }

    /// Reuse the parameters of the conquered queue, while copying the events from the source
    /// network. This requires that the configuration and the IGP state is ireused.
    ///
    /// # Safety
    /// The caller must ensure that the properties of of both the source and the conquered network
    /// queue is identical.
    pub unsafe fn reuse_queue_params(mut self, b: bool) -> Self {
        self.reuse_queue_params = b;
        self
    }

    /// Move the conquer network while cloning the required parameters from the source network into
    /// the target network.
    ///
    /// # Safety
    /// You must ensure that the physical topology of both the source and the conquered network is
    /// identical.
    pub unsafe fn conquer(self, other: Network<P, Q, GlobalOspf>) -> Network<P, Q, GlobalOspf>
    where
        Q: Clone + EventQueue<P>,
    {
        // assert that the properties are correct
        if self.reuse_igp_state && !self.reuse_config {
            panic!("Cannot reuse the IGP state but not reuse the configuration.");
        }
        if self.reuse_bgp_state && !(self.reuse_config && self.reuse_advertisements) {
            panic!(
                "Cannot reuse the BGP state but not reuse the configuration or the advertisements."
            );
        }
        if self.reuse_queue_params && !self.reuse_igp_state {
            panic!("Cannot reuse queue parameters but not reuse the IGP state.");
        }

        let mut new = other;
        let source = self.source;

        // take the values that are fast to clone
        new.stop_after = source.stop_after;
        new.skip_queue = source.skip_queue;

        // clone new.net if the configuration is different
        if !self.reuse_config {
            new.ospf.clone_from(&source.ospf);
        }

        if !self.reuse_advertisements {
            new.known_prefixes.clone_from(&source.known_prefixes);
        }

        if self.reuse_queue_params {
            new.queue = source.queue.clone_events(new.queue);
        } else {
            new.queue.clone_from(&source.queue);
        }

        // handle all internal routers
        for r in new.routers_mut() {
            let id = r.router_id();
            let r_source = source.get_router(id).unwrap();

            if !self.reuse_config {
                r.do_load_balancing = r_source.do_load_balancing;
                r.ospf.neighbors.clone_from(&r_source.ospf.neighbors);
                r.sr.clone_from(&r_source.sr);
                r.bgp.sessions.clone_from(&r_source.bgp.sessions);
                r.bgp.sessions.clone_from(&r_source.bgp.sessions);
                r.bgp.route_maps_in.clone_from(&r_source.bgp.route_maps_in);
                r.bgp
                    .route_maps_out
                    .clone_from(&r_source.bgp.route_maps_out);
            }

            if !self.reuse_igp_state {
                r.ospf.ospf_table.clone_from(&r_source.ospf.ospf_table);
            }

            if !self.reuse_bgp_state {
                r.bgp.rib_in.clone_from(&r_source.bgp.rib_in);
                r.bgp.rib.clone_from(&r_source.bgp.rib);
                r.bgp.rib_out.clone_from(&r_source.bgp.rib_out);
                r.bgp
                    .known_prefixes
                    .clone_from(&r_source.bgp.known_prefixes);
            }
        }

        new
    }
}
