// 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 functions to save the network to a file, and restore it from a file.

use std::collections::HashMap;

#[cfg(feature = "topology_zoo")]
use geoutils::Location;
use itertools::Itertools;
#[cfg(feature = "topology_zoo")]
use mapproj::{cylindrical::mer::Mer, LonLat, Projection};
use serde::{Deserialize, Serialize};
use serde_json::json;

#[cfg(feature = "topology_zoo")]
use crate::topology_zoo::TopologyZoo;
use crate::{
    config::{ConfigExpr, ConfigModifier, NetworkConfig},
    event::{BasicEventQueue, Event, EventQueue},
    network::Network,
    ospf::{LocalOspf, OspfImpl},
    policies::{FwPolicy, Policy, PolicyError},
    types::{IntoIpv4Prefix, Ipv4Prefix, NetworkError, Prefix, RouterId, ASN},
};

const JSON_FIELD_NAME_NETWORK: &str = "net";
const JSON_FIELD_NAME_CONFIG: &str = "config_nodes_routes";

type ExportConfig<P> = Vec<ConfigExpr<P>>;
type ExportRouters = Vec<(RouterId, String, ASN)>;
type ExportLinks = Vec<(RouterId, RouterId)>;
type ExportTuple<P> = (ExportConfig<P>, ExportRouters, ExportLinks);

impl<P, Q, Ospf> Network<P, Q, Ospf>
where
    P: Prefix,
    Q: EventQueue<P> + Serialize,
    Ospf: OspfImpl,
{
    /// Create a json string from the network. This string will contain both the actual network
    /// state and the configuration. In case the network state can no longer be deserialized, the
    /// configuration can be used to restore the network to a similar state.
    pub fn as_json_str(&self) -> String {
        serde_json::to_string(&json!({
            JSON_FIELD_NAME_NETWORK: serde_json::to_value(self).unwrap(),
            JSON_FIELD_NAME_CONFIG: self.as_config_json_value(),
        }))
        .unwrap()
    }

    /// Create a json string from the network. This string will only contain the configuration, not
    /// the serialized network itself. Thus, this string is significantly smaller than the string
    /// generated by `Network::as_json_str`.
    pub fn as_json_str_compact(&self) -> String {
        serde_json::to_string(&json!({
            JSON_FIELD_NAME_CONFIG: self.as_config_json_value()
        }))
        .unwrap()
    }

    /// Create a json value containing the configuration.
    fn as_config_json_value(&self) -> serde_json::Value {
        serde_json::to_value(self.as_topo_config()).unwrap()
    }
}

impl<P, Q, Ospf> Network<P, Q, Ospf>
where
    P: Prefix,
    Q: EventQueue<P>,
    Ospf: OspfImpl,
{
    /// Create a json value containing the configuration.
    fn as_topo_config(&self) -> ExportTuple<P> {
        let config = Vec::from_iter(self.get_config().unwrap().iter().cloned());
        let mut nodes: ExportRouters = self
            .routers()
            .map(|r| (r.router_id(), r.name().to_string(), r.asn()))
            .collect();
        nodes.sort_by_key(|(r, _, _)| *r);

        let links: ExportLinks = self
            .ospf
            .edges()
            .map(|e| (e.src(), e.dst()))
            .map(|(a, b)| if a < b { (a, b) } else { (b, a) })
            .unique()
            .collect();

        (config, nodes, links)
    }
}

/// A data structure to generate json files for importing into [bgpsim.github.io](bgpsim.github.io).
#[derive(Debug, Serialize)]
#[allow(clippy::type_complexity)]
pub struct WebExporter {
    net: Option<Network<Ipv4Prefix, BasicEventQueue<Ipv4Prefix>, LocalOspf>>,
    config_node_routes: ExportTuple<Ipv4Prefix>,
    pos: Option<HashMap<RouterId, Point>>,
    spec:
        Option<HashMap<RouterId, Vec<(FwPolicy<Ipv4Prefix>, Result<(), PolicyError<Ipv4Prefix>>)>>>,
    #[cfg(feature = "topology_zoo")]
    topology_zoo: Option<TopologyZoo>,
    replay: Option<Vec<(Event<Ipv4Prefix, ()>, Option<usize>)>>,
    #[serde(skip)]
    compact: bool,
}

impl WebExporter {
    /// Create a new, empty web exporter.
    pub fn new<P: Prefix, Q: EventQueue<P> + Clone, Ospf: OspfImpl>(
        net: &Network<P, Q, Ospf>,
    ) -> Self {
        // transform the queue into a basic event queue
        let mut net = net.clone();
        let mut queue: BasicEventQueue<Ipv4Prefix> = Default::default();
        while let Some(e) = net.queue.pop() {
            queue.0.push_back(e.into_ipv4_prefix())
        }
        // make OSPF local
        let net = net.into_global_ospf().unwrap().into_local_ospf().unwrap();

        // only now, swap the queue.
        let Ok(net) = net.into_ipv4_prefix(queue) else {
            unreachable!("Queue was emptied above")
        };

        // create the config
        let config_node_routes = net.as_topo_config();

        Self {
            net: Some(net),
            config_node_routes,
            pos: None,
            spec: None,
            #[cfg(feature = "topology_zoo")]
            topology_zoo: None,
            replay: None,
            compact: false,
        }
    }

    /// Remove the serialized network from the json, only keeping the configuration, the topology,
    /// and the advertised routes.
    pub fn compact(mut self) -> Self {
        self.compact = true;
        self
    }

    /// Set the specification while exporting
    pub fn spec<P: Prefix>(mut self, spec: Vec<FwPolicy<P>>) -> Self {
        // first, get the forwarding state
        let Some(net) = self.net.as_ref() else {
            self.spec = Some(
                spec.into_iter()
                    .map(FwPolicy::into_ipv4_prefix)
                    .map(|s| (s.router().unwrap(), (s, Ok(()))))
                    .into_group_map(),
            );
            return self;
        };
        let mut fw_state = net.get_forwarding_state();
        // transform the spec into the required data format.
        self.spec = Some(
            spec.into_iter()
                .map(FwPolicy::into_ipv4_prefix)
                .map(|s| (s.router().unwrap(), (s.clone(), s.check(&mut fw_state))))
                .into_group_map(),
        );
        self
    }

    /// Set the topology zoo variant that was used to create this network. Any previously set
    /// positions will be overwritten.
    #[cfg(feature = "topology_zoo")]
    pub fn topology_zoo(mut self, topo: TopologyZoo) -> Self {
        // build the position
        fn rad(x: Location) -> LonLat {
            let mut lon = x.longitude();
            let mut lat = x.latitude();
            if lon < 0.0 {
                lon += 360.0;
            }
            lon = lon * std::f64::consts::PI / 180.0;
            lat = lat * std::f64::consts::PI / 180.0;
            LonLat::new(lon, lat)
        }

        let mut geo = topo.geo_location();
        geo.retain(|_, pos| pos.latitude() != 0.0 || pos.longitude() != 0.0);
        if geo.is_empty() {
            // nothing to do, we don't have any geo information available.
            return self;
        }
        let mut pos = std::collections::HashMap::new();
        let proj = Mer::new();
        for (r, p) in geo {
            let xy = proj.proj_lonlat(&rad(p)).unwrap();
            pos.insert(
                r,
                Point {
                    x: xy.x(),
                    y: -xy.y(),
                },
            );
        }

        self.topology_zoo = Some(topo);
        self.pos = Some(pos);
        self
    }

    /// Set the positions of each router. Any position that is not fixed will be automatically
    /// determined using the spring-layout when importing into the web-app. This function will
    /// modify the existing positions accordingly, overwriting existing positions while keeping the
    /// old ones.
    pub fn set_positions(mut self, pos: HashMap<RouterId, Point>) -> Self {
        self.pos.get_or_insert(Default::default()).extend(pos);
        self
    }

    /// Replay a recording of events.
    pub fn replay<P: Prefix, T>(mut self, events: Vec<Event<P, T>>) -> Self {
        self.replay = Some(
            events
                .into_iter()
                .map(|x| (x.into_ipv4_prefix(), None))
                .collect(),
        );
        self
    }

    /// Replay a recording of events. Each event is associated with the index of the event that triggered it.
    pub fn replay_with_trigger<P: Prefix, T>(
        mut self,
        events: Vec<(Event<P, T>, Option<usize>)>,
    ) -> Self {
        self.replay = Some(
            events
                .into_iter()
                .map(|(x, id)| (x.into_ipv4_prefix(), id))
                .collect(),
        );
        self
    }

    /// Generate a json string that only contains the replay. This can be imported into
    /// bgpsim.github.io.
    pub fn replay_only<P: Prefix, T>(events: Vec<Event<P, T>>) -> String {
        let events: Vec<(Event<Ipv4Prefix, ()>, Option<usize>)> = events
            .into_iter()
            .map(|x| (x.into_ipv4_prefix(), None))
            .collect();
        serde_json::json!({
            "replay": events
        })
        .to_string()
    }

    /// Generate a json string that only contains the replay. This can be imported into
    /// bgpsim.github.io. Each event is associated with the index of the event that triggered it.
    pub fn replay_only_with_trigger<P: Prefix, T>(
        events: Vec<(Event<P, T>, Option<usize>)>,
    ) -> String {
        let events: Vec<(Event<Ipv4Prefix, ()>, Option<usize>)> = events
            .into_iter()
            .map(|(x, id)| (x.into_ipv4_prefix(), id))
            .collect();
        serde_json::json!({
            "replay": events
        })
        .to_string()
    }

    /// Create a json string representing the network.
    pub fn to_json(mut self) -> String {
        if self.compact {
            self.net = None;
        }

        serde_json::to_string(&self).unwrap()
    }
}

/// A point in 2 dimensional space, used to define the positioning of nodes.
#[derive(Debug, Serialize)]
pub struct Point {
    /// The x coordinate. The scaling is done automatically by the web-app.
    pub x: f64,
    /// The y coordinate. The scaling is done automatically by the web-app.
    pub y: f64,
}

impl<P, Q, Ospf> Network<P, Q, Ospf>
where
    P: Prefix,
    Q: EventQueue<P>,
    Ospf: OspfImpl,
    for<'a> Q: Deserialize<'a>,
{
    /// Read a json file containing the network and create the network. If the network cannot be
    /// deserialized directly, reconstruct it from the configuration that should also be part of the
    /// exported file.
    ///
    /// The `default_queue` function must return a queue in case the network cannot be deserialized
    /// directly, but it needs to be built up from the configuration. For instance, use
    /// `Default::default` for a queue `Q` like `BasicEventQueue` that implements `Default`.
    pub fn from_json_str<F>(s: &str, default_queue: F) -> Result<Self, NetworkError>
    where
        F: FnOnce() -> Q,
    {
        // first, try to deserialize the network. If that works, ignore the config
        let content: serde_json::Value = serde_json::from_str(s)?;
        if let Some(net) = content
            .get(JSON_FIELD_NAME_NETWORK)
            .and_then(|v| serde_json::from_value(v.clone()).ok())
        {
            Ok(net)
        } else {
            match content
                .get(JSON_FIELD_NAME_CONFIG)
                .and_then(|v| v.as_array())
            {
                Some(v) if v.len() == 3 => Self::from_config_nodes_routes(
                    v[0].clone(),
                    v[1].clone(),
                    v[2].clone(),
                    default_queue,
                ),
                _ => Err(serde_json::from_str::<ConfigNodeRoutes>(s).unwrap_err())?,
            }
        }
    }
}

impl<P, Q, Ospf> Network<P, Q, Ospf>
where
    P: Prefix,
    Q: EventQueue<P>,
    Ospf: OspfImpl,
{
    /// Deserialize the json structure containing configuration, nodes and routes.
    fn from_config_nodes_routes<F>(
        config: serde_json::Value,
        nodes: serde_json::Value,
        links: serde_json::Value,
        default_queue: F,
    ) -> Result<Self, NetworkError>
    where
        F: FnOnce() -> Q,
    {
        let config: Vec<ConfigExpr<P>> = serde_json::from_value(config)?;
        let nodes: ExportRouters = serde_json::from_value(nodes)?;
        let links: ExportLinks = serde_json::from_value(links)?;
        let mut nodes_lut: HashMap<RouterId, RouterId> = HashMap::new();
        let mut net = Network::new(default_queue());
        // add all nodes and create the lut
        for (id, name, asn) in nodes.into_iter() {
            let new_id = net.add_router(name, asn);
            nodes_lut.insert(id, new_id);
        }
        // create the function to lookup nodes
        let node = |id: RouterId| {
            nodes_lut
                .get(&id)
                .copied()
                .ok_or(NetworkError::DeviceNotFound(id))
        };
        let links = links
            .into_iter()
            .map(|(a, b)| Ok::<_, NetworkError>((node(a)?, node(b)?)))
            .collect::<Result<Vec<_>, _>>()?;
        net.add_links_from(links)?;
        // apply all configurations

        for expr in config.iter() {
            let expr = match expr.clone() {
                ConfigExpr::IgpLinkWeight {
                    source,
                    target,
                    weight,
                } => ConfigExpr::IgpLinkWeight {
                    source: node(source)?,
                    target: node(target)?,
                    weight,
                },
                ConfigExpr::OspfArea {
                    source,
                    target,
                    area,
                } => ConfigExpr::OspfArea {
                    source: node(source)?,
                    target: node(target)?,
                    area,
                },
                ConfigExpr::BgpSession {
                    source,
                    target,
                    target_is_client,
                } => ConfigExpr::BgpSession {
                    source: node(source)?,
                    target: node(target)?,
                    target_is_client,
                },
                ConfigExpr::BgpRouteMap {
                    router,
                    neighbor,
                    direction,
                    map,
                } => ConfigExpr::BgpRouteMap {
                    router: node(router)?,
                    neighbor: node(neighbor)?,
                    direction,
                    map,
                },
                ConfigExpr::StaticRoute {
                    router,
                    prefix,
                    target,
                } => ConfigExpr::StaticRoute {
                    router: node(router)?,
                    prefix,
                    target,
                },
                ConfigExpr::LoadBalancing { router } => ConfigExpr::LoadBalancing {
                    router: node(router)?,
                },
                ConfigExpr::AdvertiseRoute {
                    router,
                    prefix,
                    as_path,
                    med,
                    community,
                } => ConfigExpr::AdvertiseRoute {
                    router: node(router)?,
                    prefix,
                    as_path,
                    med,
                    community,
                },
            };
            net.apply_modifier(&ConfigModifier::Insert(expr))?;
        }
        Ok(net)
    }
}

/// Dummy struct that allows us to create meaningful error messages
#[derive(Debug, Deserialize)]
struct ConfigNodeRoutes {
    #[allow(dead_code)]
    config_nodes_routes: (serde_json::Value, serde_json::Value, serde_json::Value),
}
