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

//! Export an external router into files for [ExaBGP](https://github.com/Exa-Networks/exabgp).

use std::{
    collections::{BTreeMap, BTreeSet},
    net::Ipv4Addr,
    time::Duration,
};

use crate::{
    bgp::BgpRoute,
    network::Network,
    ospf::OspfImpl,
    types::{Ipv4Prefix, Prefix, PrefixMap, RouterId, ASN},
};

use super::{Addressor, CfgGen, ExportError};

use ipnet::Ipv4Net;
use itertools::Itertools;
use maplit::btreemap;

/// The python preamble of the runner script (to import time and sleep for 5 seconds)
pub const RUNNER_PREAMBLE: &str =
    "#!/usr/bin/env python3\n\nimport sys\nimport time\n\n\ntime.sleep(5)\n\n";

/// The python preamble of the runner script (to import time and sleep for 5 seconds)
pub const RUNNER_POSTAMBLE: &str = "\nwhile True:\n    time.sleep(1)\n";

/// Config generator for [ExaBGP](https://github.com/Exa-Networks/exabgp)
///
/// This generator works differently. Instead of giving the configuration from one single time
/// instance, it tries to give the configuration for an entire sequence. When calling
/// `generate_config`, it will generate the configuration file for exabgp, which will create the
/// necessary sessions. However, when calling `advertise_route`, or `withdraw_route`, the function
/// will return a python script that loops forever, and advertises or withdraws routes accordingly.
///
/// This structure will keep a history of all routes, along with the time at which they should be
/// advertised or withdrawn. When calling either `advertise_route` or `withdraw_route`, this will
/// push a new entry for this route into the history, at the time set by calling
/// `step_time`.
///
/// All events are triggered once, and the script will go into an infinite loop.
///
/// ## Configuation
///
/// `ExaBgpCfgGen` implements [`ExternalCfgGen`]. When calling [`ExternalCfgGen::generate_config`],
/// then the configuration block is created only for those sessions of the specific router. They
/// look as follows:
///
/// ```text
/// neighbor 10.192.0.1 {
///     router-id 20.0.0.1;
///     local-address 10.255.0.1;
///     local-as 100;
///     peer-as 65535;
///     hold-time 180;
///     family { ipv4 unicast; }
///     capability { route-refresh; }
/// }
/// ```
///
/// ## Python Runner
///
/// Further, when calling [`ExternalCfgGen::advertise_route`], [`ExternalCfgGen::withdraw_route`],
/// or [`ExaBgpCfgGen::generate_script`], then a python script is generated that looks as follows:
///
/// ```py
/// #!/usr/bin/env python3
///
/// import sys
/// import time
///
///
/// time.sleep(5)
///
/// sys.stdout.write("neighbor 10.192.0.1 announce route 100.0.0.0/16 next-hop self as-path [100]\n")
/// sys.stdout.write("neighbor 10.192.0.1 announce route 100.1.0.0/16 next-hop self as-path [100, 200, 300]\n")
/// sys.stdout.flush()
/// time.sleep(10)
/// sys.stdout.write("neighbor 10.192.0.1 withdraw route 100.0.0.0/16\n")
/// sys.stdout.flush()
/// time.sleep(10)
/// sys.stdout.write("neighbor 10.192.0.1 announce route 100.1.0.0/16 next-hop self as-path [100, 300]\n")
/// sys.stdout.flush()
///
/// while True:
///     time.sleep(1)
/// ```
///
/// ## Example
///
/// The two files above are generated by the following code:
///
/// ```
/// use std::time::Duration;
/// use bgpsim::prelude::*;
/// use bgpsim::types::SimplePrefix as P;
/// use bgpsim::export::{DefaultAddressorBuilder, ExternalCfgGen, ExaBgpCfgGen};
/// # use pretty_assertions::assert_eq;
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // create the network and get the external router
/// let mut net = {
///     // ...
/// #   use bgpsim::builder::*;
/// #   let mut net: Network<_, _, GlobalOspf> = Network::new(BasicEventQueue::<P>::new());
/// #   net.build_topology(65500, CompleteGraph(1)).unwrap();
/// #   let router = net.add_router("external_router", ASN(100));
/// #   net.internal_indices().detach().for_each(|r| net.add_link(r, router).unwrap());
/// #   net.build_ibgp_full_mesh()?;
/// #   net.build_ebgp_sessions()?;
/// #   net.build_link_weights(1.0)?;
/// #   net
/// };
/// let router = net.get_router_id("external_router")?;
///
/// // Create some advertisements
/// net.advertise_external_route(router, P::from(0), [100], None, None)?;
/// net.advertise_external_route(router, P::from(1), [100, 200, 300], None, None)?;
///
/// // create the addressor
/// let mut addressor = DefaultAddressorBuilder::default().build(&net)?;
///
/// // create the config generator
/// let mut cfg = ExaBgpCfgGen::new(&net, router)?;
///
/// // generate the configuration
/// assert_eq!(
///     cfg.generate_config(&net, &mut addressor)?,
///     "\
/// neighbor 1.192.0.1 {
///     router-id 2.0.0.1;
///     local-address 1.192.0.2;
///     local-as 100;
///     peer-as 65500;
///     family { ipv4 unicast; }
///     capability { route-refresh; }
/// }"
/// );
///
/// // create a script that withdraws the route for prefix 0 after 10 seconds, and changes the AS
/// // path of prefix 1 after 20 seconds
/// cfg.step_time(Duration::from_secs(10));
/// cfg.withdraw_route(&net, &mut addressor, P::from(0))?;
/// cfg.step_time(Duration::from_secs(10));
/// cfg.advertise_route(
///     &net,
///     &mut addressor,
///     &BgpRoute::new(router, P::from(1), [100, 300], None, None)
/// )?;
///
/// // generate the script
/// assert_eq!(
///     cfg.generate_script(&mut addressor)?,
///     "\
/// #!/usr/bin/env python3
///
/// import sys
/// import time
///
///
/// time.sleep(5)
///
/// sys.stdout.write(\"neighbor 1.192.0.1 announce route 100.0.0.0/24 next-hop self as-path [100]\\n\")
/// sys.stdout.write(\"neighbor 1.192.0.1 announce route 100.0.1.0/24 next-hop self as-path [100, 200, 300]\\n\")
/// sys.stdout.flush()
/// time.sleep(10)
/// sys.stdout.write(\"neighbor 1.192.0.1 withdraw route 100.0.0.0/24\\n\")
/// sys.stdout.flush()
/// time.sleep(10)
/// sys.stdout.write(\"neighbor 1.192.0.1 announce route 100.0.1.0/24 next-hop self as-path [100, 300]\\n\")
/// sys.stdout.flush()
///
/// while True:
///     time.sleep(1)
/// "
/// );
///
/// # Ok(()) }
/// ```
#[derive(Debug)]
pub struct ExaBgpCfgGen<P: Prefix> {
    router: RouterId,
    asn: ASN,
    routes: BTreeMap<P, BTreeMap<Duration, Option<BgpRoute<P>>>>,
    neighbors: BTreeSet<RouterId>,
    current_time: Duration,
}

impl<P: Prefix> ExaBgpCfgGen<P> {
    /// Create a new instance of the ExaBGP config generator. This will initialize all
    /// routes. Further, it will
    pub fn new<Q, Ospf: OspfImpl>(
        net: &Network<P, Q, Ospf>,
        router: RouterId,
    ) -> Result<Self, ExportError> {
        let r = net.get_device(router)?.external_or_err()?;
        Ok(Self {
            router,
            asn: r.asn(),
            routes: r
                .active_routes
                .iter()
                .map(|(p, r)| (*p, btreemap! {Duration::ZERO => Some(r.clone())}))
                .collect(),
            neighbors: r.neighbors.iter().copied().collect(),
            current_time: Duration::ZERO,
        })
    }

    /// Increase the `current_time` by the given amount.
    ///
    /// After creating a new instance of `ExaBgpCfgGen`, the `current_time` will be set to 0.
    pub fn step_time(&mut self, step: Duration) {
        self.current_time += step;
    }

    /// Generate the python script that loops over the history of routes, and replays that over and
    /// over again.
    pub fn generate_script<A: Addressor<P>>(
        &self,
        addressor: &mut A,
    ) -> Result<String, ExportError> {
        let script = String::from(RUNNER_PREAMBLE);

        Ok(script + &self.generate_script_no_loop(addressor)?)
    }

    /// Generate all python command lines to advertise or withdraw the routes. This will create a
    /// vector of strings for each time step, and return all of these time steps along with the
    /// duration when they should be triggered.
    pub fn generate_lines<A: Addressor<P>>(
        &self,
        addressor: &mut A,
    ) -> Result<Vec<(Vec<String>, Duration)>, ExportError> {
        let neighbors = self
            .neighbors
            .iter()
            .map(|x| addressor.iface_address(*x, self.router))
            .collect::<Result<Vec<Ipv4Addr>, ExportError>>()?
            .into_iter()
            .map(|x| format!("neighbor {x}"))
            .join(", ");

        let mut result = Vec::new();

        let mut times_routes: BTreeMap<_, Vec<_>> = Default::default();
        for (p, routes) in self.routes.iter() {
            for (time, route) in routes.iter() {
                times_routes
                    .entry(*time)
                    .or_default()
                    .push((*p, route.as_ref()));
            }
        }

        for (time, routes) in times_routes {
            let mut ads: Vec<String> = Vec::new();
            for (p, r) in routes {
                for net in addressor.prefix(p)? {
                    if let Some(r) = r {
                        let r = r.clone().with_prefix(Ipv4Prefix::from(net));
                        ads.push(format!(
                            "sys.stdout.write(\"{neighbors} {}\\n\")",
                            announce_route(&r)
                        ))
                    } else {
                        ads.push(format!(
                            "sys.stdout.write(\"{neighbors} {}\\n\")",
                            withdraw_route(net)
                        ))
                    }
                }
            }
            result.push((ads, time));
        }

        Ok(result)
    }

    /// Generate the python script that does not loop, but trigger the events once. The header of
    /// the script is not generated!
    fn generate_script_no_loop<A: Addressor<P>>(
        &self,
        addressor: &mut A,
    ) -> Result<String, ExportError> {
        let lines = self.generate_lines(addressor)?;

        let mut script = String::new();

        let mut current_time = Duration::ZERO;
        for (routes, time) in lines {
            if !time.is_zero() {
                script.push_str(&format!(
                    "time.sleep({})\n",
                    (time - current_time).as_secs_f64()
                ));
            }
            current_time = time;
            for route in routes {
                script.push_str(&route);
                script.push('\n');
            }
            script.push_str("sys.stdout.flush()\n");
        }

        script.push_str(RUNNER_POSTAMBLE);

        Ok(script)
    }

    /// Generate the configuration for a single neighbor
    fn generate_neighbor_cfg<A: Addressor<P>, Q, Ospf: OspfImpl>(
        &self,
        net: &Network<P, Q, Ospf>,
        addressor: &mut A,
        neighbor: RouterId,
    ) -> Result<String, ExportError> {
        let asn = net.get_device(neighbor)?.asn();
        Ok(format!(
            "\
neighbor {} {{
    router-id {};
    local-address {};
    local-as {};
    peer-as {};
    family {{ ipv4 unicast; }}
    capability {{ route-refresh; }}
}}",
            addressor.iface_address(neighbor, self.router)?,
            addressor.router_address(self.router)?,
            addressor.iface_address(self.router, neighbor)?,
            self.asn.0,
            asn.0,
        ))
    }

    /// Function to get all neighbors of that external router.
    pub fn neighbors(&self) -> &BTreeSet<RouterId> {
        &self.neighbors
    }
}

/// Get the text to announce a route.
pub fn announce_route<P: Prefix>(route: &BgpRoute<P>) -> String {
    let prefix: Ipv4Net = route.prefix.into();
    format!(
        "announce route {prefix} next-hop self as-path [{}]{}{}",
        route.as_path.iter().map(|x| x.0).join(", "),
        if let Some(med) = route.med {
            format!(" metric {med}")
        } else {
            String::new()
        },
        if route.community.is_empty() {
            String::new()
        } else {
            format!(
                " extended-community [{}]",
                route
                    .community
                    .iter()
                    .map(|x| format!("{}:{}", 65535, x)) // TODO replace with proper community
                    .join(", ")
            )
        },
    )
}

/// Get the text to withdraw a route.
pub fn withdraw_route(prefix: Ipv4Net) -> String {
    format!("withdraw route {prefix}")
}

impl<P: Prefix, A: Addressor<P>, Q, Ospf: OspfImpl> ExternalCfgGen<P, Q, Ospf, A>
    for ExaBgpCfgGen<P>
{
    fn generate_config(
        &mut self,
        net: &Network<P, Q, Ospf>,
        addressor: &mut A,
    ) -> Result<String, ExportError> {
        Ok(self
            .neighbors
            .iter()
            .map(|x| self.generate_neighbor_cfg(net, addressor, *x))
            .collect::<Result<Vec<String>, ExportError>>()?
            .into_iter()
            .join("\n"))
    }

    fn advertise_route(
        &mut self,
        _net: &Network<P, Q, Ospf>,
        addressor: &mut A,
        route: &BgpRoute<P>,
    ) -> Result<String, ExportError> {
        self.routes
            .entry(route.prefix)
            .or_default()
            .insert(self.current_time, Some(route.clone()));
        self.generate_script(addressor)
    }

    fn withdraw_route(
        &mut self,
        _net: &Network<P, Q, Ospf>,
        addressor: &mut A,
        prefix: P,
    ) -> Result<String, ExportError> {
        self.routes
            .entry(prefix)
            .or_default()
            .insert(self.current_time, None);
        self.generate_script(addressor)
    }

    fn establish_ebgp_session(
        &mut self,
        net: &Network<P, Q, Ospf>,
        addressor: &mut A,
        neighbor: RouterId,
    ) -> Result<String, ExportError> {
        self.neighbors.insert(neighbor);
        self.generate_config(net, addressor)
    }

    fn teardown_ebgp_session(
        &mut self,
        net: &Network<P, Q, Ospf>,
        addressor: &mut A,
        neighbor: RouterId,
    ) -> Result<String, ExportError> {
        self.neighbors.remove(&neighbor);
        self.generate_config(net, addressor)
    }
}
