//! What a policy function is given: the route and the peer.
//!
//! # The route
//!
//! | In a script | |
//! |---|---|
//! | `route.prefix` | `"203.0.113.0/24"` |
//! | `route.prefix_len` | `24` |
//! | `route.within("203.0.112.0/22")` | whether the prefix is that one or inside it |
//! | `route.family` | `"ipv4"` or `"ipv6"` |
//! | `route.as_path` | `[64602, 3356, 15169]`, the neighbor's AS first |
//! | `route.as_path_len` | as route selection counts it: a set is one |
//! | `route.origin_as`, `route.neighbor_as` | the last and the first AS of the path; 0 if there is none |
//! | `route.path_has(3356)` | whether the AS is anywhere on the path |
//! | `route.origin` | `"igp"`, `"egp"` or `"incomplete"` |
//! | `route.next_hop` | `"192.0.2.1"` |
//! | `route.med`, `route.local_pref` | the number, or `()` if the route has none; both can be set, and unset with `()` |
//! | `route.communities` | `["3356:3", "3356:2012"]` |
//! | `route.has_community("3356:3")` | also `(3356, 3)`, and `"3356:*"` for any of that AS |
//! | `route.add_community("64601:100")` | also `(64601, 100)` |
//! | `route.remove_community("3356:*")` | all that match; `"*:*"` for all there are |
//! | `route.large_communities`, `route.has_large_community("64601:1:*")`, `route.add_large_community(..)`, `route.remove_large_community(..)` | the same for large communities (RFC 8092) |
//! | `route.prepend(2)` | on the way out: our AS that many more times on the path |
//!
//! # The peer
//!
//! `peer.name` (as configured), `peer.asn`, `peer.address`, `peer.role`
//! (`"customer"`, `"provider"`, `"peer"`, `"route_server"`,
//! `"route_server_client"`, or `""` if none is configured), `peer.ibgp`,
//! and `peer.local_asn`, which is our own.

use std::net::IpAddr;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};

use bgp_packet::community::{format_large, format_standard};
use bgp_packet::message::{
    AsPathAttribute, AsPathSegment, CommunitiesPathAttribute, LargeCommunitiesPathAttribute,
    LargeCommunity, LocalPrefPathAttribute, MultiExitDiscPathAttribute, NlriNextHop,
    OriginPathAttribute, PathAttribute,
};
use bgp_primitives::constants::IpVersion;
use bgp_primitives::ip_prefix::IpPrefix;
use rhai::{Array, Dynamic, Engine, EvalAltResult, INT, ImmutableString};

use crate::{Direction, Subject};

type Fallible<T> = Result<T, Box<EvalAltResult>>;

/// A peer, as a script sees it. Made once for a session.
#[derive(Debug, Clone)]
pub struct Peer {
    name: ImmutableString,
    address: ImmutableString,
    role: ImmutableString,
    asn: u32,
    local_asn: u32,
    ibgp: bool,
}

impl Peer {
    /// `role` is the peer's role by the name scripts know it under, or
    /// empty.
    pub fn new(name: &str, address: IpAddr, asn: u32, local_asn: u32, role: &str) -> Self {
        Self {
            name: name.into(),
            address: address.to_string().into(),
            role: role.into(),
            asn,
            local_asn,
            ibgp: asn == local_asn,
        }
    }
}

/// The route a function is looking at, and what it has done to it.
#[derive(Debug)]
struct State {
    prefix: IpPrefix,
    direction: Direction,
    local_asn: u32,
    ibgp: bool,
    /// The AS numbers of the path in order, sets and all.
    as_path: Vec<u32>,
    as_path_len: usize,
    origin_as: u32,
    neighbor_as: u32,
    origin: OriginPathAttribute,
    next_hop: Option<NlriNextHop>,
    communities: Vec<(u16, u16)>,
    large_communities: Vec<LargeCommunity>,
    local_pref: Option<u32>,
    med: Option<u32>,
    /// How many more times our AS is to go on the path.
    prepend: usize,
    saw_prefix: bool,
    /// Which of the above the function changed.
    changed: Changed,
}

#[derive(Debug, Default, Clone, Copy)]
struct Changed {
    communities: bool,
    large_communities: bool,
    local_pref: bool,
    med: bool,
}

/// The handle a script holds. The function is given a copy and changes
/// what it points to; the caller reads that afterwards.
#[derive(Debug, Clone)]
pub(crate) struct Route(Arc<Mutex<State>>);

impl Route {
    pub(crate) fn new(subject: &Subject<'_>, peer: &Peer) -> Self {
        let mut state = State {
            prefix: *subject.prefix,
            direction: subject.direction,
            local_asn: peer.local_asn,
            ibgp: peer.ibgp,
            as_path: Vec::new(),
            as_path_len: 0,
            origin_as: 0,
            neighbor_as: 0,
            origin: OriginPathAttribute::INCOMPLETE,
            next_hop: subject.next_hop,
            communities: Vec::new(),
            large_communities: Vec::new(),
            local_pref: None,
            med: None,
            prepend: 0,
            saw_prefix: false,
            changed: Changed::default(),
        };
        for attribute in subject.attributes {
            match attribute {
                PathAttribute::Origin(origin) => state.origin = *origin,
                PathAttribute::AsPath(path) => {
                    state.as_path_len = path.hops();
                    let first = path.segments.first().filter(|segment| segment.ordered);
                    let last = path.segments.last().filter(|segment| segment.ordered);
                    state.neighbor_as = first.and_then(|s| s.path.first()).copied().unwrap_or(0);
                    state.origin_as = last.and_then(|s| s.path.last()).copied().unwrap_or(0);
                    state.as_path = path
                        .segments
                        .iter()
                        .flat_map(|segment| segment.path.iter().copied())
                        .collect();
                }
                PathAttribute::NextHop(next_hop) => {
                    state.next_hop = Some(NlriNextHop::Ipv4(next_hop.0));
                }
                PathAttribute::MultiExitDisc(med) => state.med = Some(med.0),
                PathAttribute::LocalPref(local_pref) => state.local_pref = Some(local_pref.0),
                PathAttribute::Communities(communities) => {
                    state.communities.clone_from(&communities.0);
                }
                PathAttribute::LargeCommunities(large) => {
                    state.large_communities.clone_from(&large.communities);
                }
                _ => {}
            }
        }
        Self(Arc::new(Mutex::new(state)))
    }

    fn state(&self) -> MutexGuard<'_, State> {
        self.0.lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// What the function left: `attributes` as changed, if it changed
    /// any, and whether it asked for the prefix.
    pub(crate) fn finish(
        &self,
        attributes: &[PathAttribute],
    ) -> (Option<Vec<PathAttribute>>, bool) {
        let state = self.state();
        let changed = state.changed;
        let prepend = state.prepend > 0 && !state.ibgp;
        let any = changed.communities
            || changed.large_communities
            || changed.local_pref
            || changed.med
            || prepend;
        if !any {
            return (None, state.saw_prefix);
        }
        let mut attributes = attributes.to_vec();
        if changed.communities {
            let communities = (!state.communities.is_empty()).then(|| {
                PathAttribute::Communities(CommunitiesPathAttribute(state.communities.clone()))
            });
            put(
                &mut attributes,
                bgp_packet::message::attr_type::COMMUNITIES,
                communities,
            );
        }
        if changed.large_communities {
            let large = (!state.large_communities.is_empty()).then(|| {
                PathAttribute::LargeCommunities(LargeCommunitiesPathAttribute {
                    communities: state.large_communities.clone(),
                })
            });
            put(
                &mut attributes,
                bgp_packet::message::attr_type::LARGE_COMMUNITIES,
                large,
            );
        }
        if changed.local_pref {
            let local_pref = state
                .local_pref
                .map(|value| PathAttribute::LocalPref(LocalPrefPathAttribute(value)));
            put(
                &mut attributes,
                bgp_packet::message::attr_type::LOCAL_PREF,
                local_pref,
            );
        }
        if changed.med {
            let med = state
                .med
                .map(|value| PathAttribute::MultiExitDisc(MultiExitDiscPathAttribute(value)));
            put(
                &mut attributes,
                bgp_packet::message::attr_type::MULTI_EXIT_DISC,
                med,
            );
        }
        if prepend {
            let ours = std::iter::repeat_n(state.local_asn, state.prepend);
            for attribute in &mut attributes {
                if let PathAttribute::AsPath(path) = attribute {
                    prepend_to(path, ours);
                    break;
                }
            }
        }
        (Some(attributes), state.saw_prefix)
    }
}

/// Puts `ours` at the front of `path`.
fn prepend_to(path: &mut AsPathAttribute, ours: impl Iterator<Item = u32>) {
    match path.segments.first_mut() {
        Some(first) if first.ordered => {
            first.path.splice(0..0, ours);
        }
        _ => path.segments.insert(
            0,
            AsPathSegment {
                ordered: true,
                path: ours.collect(),
            },
        ),
    }
}

/// Makes the attribute of type `type_code` in `attributes` be `with`:
/// replaced, added in the order of the type codes, or taken out.
fn put(attributes: &mut Vec<PathAttribute>, type_code: u8, with: Option<PathAttribute>) {
    let at = attributes
        .iter()
        .position(|attribute| attribute.type_code() == type_code);
    match (at, with) {
        (Some(at), Some(with)) => attributes[at] = with,
        (Some(at), None) => {
            attributes.remove(at);
        }
        (None, Some(with)) => {
            let after = attributes
                .iter()
                .position(|attribute| attribute.type_code() > type_code)
                .unwrap_or(attributes.len());
            attributes.insert(after, with);
        }
        (None, None) => {}
    }
}

fn wrong(what: impl Into<String>) -> Box<EvalAltResult> {
    what.into().into()
}

/// A number a script gave, as one that fits where it is going.
fn number<T: TryFrom<INT>>(value: INT, what: &str) -> Fallible<T> {
    T::try_from(value).map_err(|_| wrong(format!("{value} is not {what}")))
}

/// One part of a community as a script writes it: a number, or `*` for
/// any.
fn part<T: std::str::FromStr>(text: &str, whole: &str) -> Fallible<Option<T>> {
    if text == "*" {
        return Ok(None);
    }
    text.parse()
        .map(Some)
        .map_err(|_| wrong(format!("{whole:?} is not a community")))
}

type Pattern = (Option<u16>, Option<u16>);
type LargePattern = (Option<u32>, Option<u32>, Option<u32>);

fn pattern(text: &str) -> Fallible<Pattern> {
    let (high, low) = text.split_once(':').ok_or_else(|| {
        wrong(format!(
            "{text:?} is not a community: it is written 64601:100"
        ))
    })?;
    Ok((part(high, text)?, part(low, text)?))
}

fn large_pattern(text: &str) -> Fallible<LargePattern> {
    let mut parts = text.splitn(3, ':');
    let mut next = || {
        parts.next().ok_or_else(|| {
            wrong(format!(
                "{text:?} is not a large community: it is written 64601:1:2"
            ))
        })
    };
    let (admin, first, second) = (next()?, next()?, next()?);
    Ok((part(admin, text)?, part(first, text)?, part(second, text)?))
}

/// A pattern with nothing left open: a community to add.
fn exact((high, low): Pattern, text: &str) -> Fallible<(u16, u16)> {
    high.zip(low).ok_or_else(|| {
        wrong(format!(
            "{text:?} cannot be added: it stands for more than one"
        ))
    })
}

fn exact_large((admin, first, second): LargePattern, text: &str) -> Fallible<LargeCommunity> {
    match (admin, first, second) {
        (Some(global_admin), Some(data_1), Some(data_2)) => Ok(LargeCommunity {
            global_admin,
            data_1,
            data_2,
        }),
        _ => Err(wrong(format!(
            "{text:?} cannot be added: it stands for more than one"
        ))),
    }
}

fn fits<T: PartialEq>(wanted: Option<T>, has: T) -> bool {
    wanted.is_none_or(|wanted| wanted == has)
}

fn matches((high, low): Pattern, community: (u16, u16)) -> bool {
    fits(high, community.0) && fits(low, community.1)
}

fn matches_large((admin, first, second): LargePattern, community: &LargeCommunity) -> bool {
    fits(admin, community.global_admin)
        && fits(first, community.data_1)
        && fits(second, community.data_2)
}

impl State {
    fn add_community(&mut self, community: (u16, u16)) {
        if !self.communities.contains(&community) {
            self.communities.push(community);
            self.changed.communities = true;
        }
    }

    fn remove_communities(&mut self, pattern: Pattern) {
        let before = self.communities.len();
        self.communities
            .retain(|community| !matches(pattern, *community));
        self.changed.communities |= self.communities.len() != before;
    }

    fn set_local_pref(&mut self, value: Option<u32>) {
        self.changed.local_pref |= self.local_pref != value;
        self.local_pref = value;
    }

    fn set_med(&mut self, value: Option<u32>) {
        self.changed.med |= self.med != value;
        self.med = value;
    }
}

fn optional(value: Option<u32>) -> Dynamic {
    value.map_or(Dynamic::UNIT, |value| Dynamic::from(INT::from(value)))
}

/// Tells `engine` what a route and a peer are, and what can be asked of
/// them.
pub(crate) fn register(engine: &mut Engine) {
    engine.register_type_with_name::<Route>("Route");
    engine.register_type_with_name::<Peer>("Peer");

    // The peer.
    engine.register_get("name", |peer: &mut Peer| peer.name.clone());
    engine.register_get("address", |peer: &mut Peer| peer.address.clone());
    engine.register_get("role", |peer: &mut Peer| peer.role.clone());
    engine.register_get("asn", |peer: &mut Peer| INT::from(peer.asn));
    engine.register_get("local_asn", |peer: &mut Peer| INT::from(peer.local_asn));
    engine.register_get("ibgp", |peer: &mut Peer| peer.ibgp);

    // The prefix. Asking for it is noted: the answer to a function that
    // did not is good for every prefix.
    engine.register_get("prefix", |route: &mut Route| {
        let mut state = route.state();
        state.saw_prefix = true;
        state.prefix.to_string()
    });
    engine.register_get("prefix_len", |route: &mut Route| {
        let mut state = route.state();
        state.saw_prefix = true;
        INT::from(state.prefix.length())
    });
    engine.register_fn(
        "within",
        |route: &mut Route, around: &str| -> Fallible<bool> {
            let around = IpPrefix::try_from(around)
                .map_err(|e| wrong(format!("{around:?} is not a prefix: {e}")))?;
            let mut state = route.state();
            state.saw_prefix = true;
            Ok(around.contains(&state.prefix))
        },
    );
    // Every route of an UPDATE's one batch is of the same family.
    engine.register_get("family", |route: &mut Route| {
        match route.state().prefix.address_family() {
            IpVersion::Ipv4 => "ipv4",
            IpVersion::Ipv6 => "ipv6",
        }
    });

    // The path.
    engine.register_get("as_path", |route: &mut Route| -> Array {
        route
            .state()
            .as_path
            .iter()
            .map(|asn| Dynamic::from(INT::from(*asn)))
            .collect()
    });
    engine.register_get("as_path_len", |route: &mut Route| {
        route.state().as_path_len as INT
    });
    engine.register_get("origin_as", |route: &mut Route| {
        INT::from(route.state().origin_as)
    });
    engine.register_get("neighbor_as", |route: &mut Route| {
        INT::from(route.state().neighbor_as)
    });
    engine.register_fn("path_has", |route: &mut Route, asn: INT| {
        u32::try_from(asn).is_ok_and(|asn| route.state().as_path.contains(&asn))
    });
    engine.register_fn("prepend", |route: &mut Route, times: INT| -> Fallible<()> {
        let mut state = route.state();
        if state.direction == Direction::Import {
            return Err(wrong(
                "prepend is for routes on their way out; to make a route less preferred on \
                 its way in, set its local_pref",
            ));
        }
        // More than this is a mistake, and makes the route no worse.
        state.prepend = (0..=16)
            .contains(&times)
            .then_some(times as usize)
            .ok_or_else(|| wrong(format!("prepend({times}): from 0 to 16 times")))?;
        Ok(())
    });

    // The other attributes that are looked at.
    engine.register_get("origin", |route: &mut Route| match route.state().origin {
        OriginPathAttribute::IGP => "igp",
        OriginPathAttribute::EGP => "egp",
        OriginPathAttribute::INCOMPLETE => "incomplete",
    });
    engine.register_get("next_hop", |route: &mut Route| {
        route
            .state()
            .next_hop
            .map_or_else(String::new, |next_hop| match next_hop {
                NlriNextHop::Ipv4(address) => address.to_string(),
                NlriNextHop::Ipv6(address) => address.to_string(),
                NlriNextHop::Ipv6WithLl { global, .. } => global.to_string(),
            })
    });

    // Those that are also set.
    engine.register_get("local_pref", |route: &mut Route| {
        optional(route.state().local_pref)
    });
    engine.register_set(
        "local_pref",
        |route: &mut Route, value: INT| -> Fallible<()> {
            let value = number(value, "a LOCAL_PREF")?;
            route.state().set_local_pref(Some(value));
            Ok(())
        },
    );
    engine.register_set("local_pref", |route: &mut Route, (): ()| {
        route.state().set_local_pref(None);
    });
    engine.register_get("med", |route: &mut Route| optional(route.state().med));
    engine.register_set("med", |route: &mut Route, value: INT| -> Fallible<()> {
        let value = number(value, "a MULTI_EXIT_DISC")?;
        route.state().set_med(Some(value));
        Ok(())
    });
    engine.register_set("med", |route: &mut Route, (): ()| {
        route.state().set_med(None);
    });

    // Communities.
    engine.register_get("communities", |route: &mut Route| -> Array {
        route
            .state()
            .communities
            .iter()
            .map(|community| Dynamic::from(format_standard(*community)))
            .collect()
    });
    engine.register_fn(
        "has_community",
        |route: &mut Route, text: &str| -> Fallible<bool> {
            let wanted = pattern(text)?;
            Ok(route
                .state()
                .communities
                .iter()
                .any(|community| matches(wanted, *community)))
        },
    );
    engine.register_fn(
        "has_community",
        |route: &mut Route, high: INT, low: INT| -> Fallible<bool> {
            let wanted = (
                number(high, "half a community")?,
                number(low, "half a community")?,
            );
            Ok(route.state().communities.contains(&wanted))
        },
    );
    engine.register_fn(
        "add_community",
        |route: &mut Route, text: &str| -> Fallible<()> {
            let community = exact(pattern(text)?, text)?;
            route.state().add_community(community);
            Ok(())
        },
    );
    engine.register_fn(
        "add_community",
        |route: &mut Route, high: INT, low: INT| -> Fallible<()> {
            let community = (
                number(high, "half a community")?,
                number(low, "half a community")?,
            );
            route.state().add_community(community);
            Ok(())
        },
    );
    engine.register_fn(
        "remove_community",
        |route: &mut Route, text: &str| -> Fallible<()> {
            let unwanted = pattern(text)?;
            route.state().remove_communities(unwanted);
            Ok(())
        },
    );
    engine.register_fn(
        "remove_community",
        |route: &mut Route, high: INT, low: INT| -> Fallible<()> {
            let unwanted = (
                Some(number(high, "half a community")?),
                Some(number(low, "half a community")?),
            );
            route.state().remove_communities(unwanted);
            Ok(())
        },
    );

    // Large communities.
    engine.register_get("large_communities", |route: &mut Route| -> Array {
        route
            .state()
            .large_communities
            .iter()
            .map(|community| Dynamic::from(format_large(community)))
            .collect()
    });
    engine.register_fn(
        "has_large_community",
        |route: &mut Route, text: &str| -> Fallible<bool> {
            let wanted = large_pattern(text)?;
            Ok(route
                .state()
                .large_communities
                .iter()
                .any(|community| matches_large(wanted, community)))
        },
    );
    engine.register_fn(
        "add_large_community",
        |route: &mut Route, text: &str| -> Fallible<()> {
            let community = exact_large(large_pattern(text)?, text)?;
            let mut state = route.state();
            if !state.large_communities.contains(&community) {
                state.large_communities.push(community);
                state.changed.large_communities = true;
            }
            Ok(())
        },
    );
    engine.register_fn(
        "remove_large_community",
        |route: &mut Route, text: &str| -> Fallible<()> {
            let unwanted = large_pattern(text)?;
            let mut state = route.state();
            let before = state.large_communities.len();
            state
                .large_communities
                .retain(|community| !matches_large(unwanted, community));
            state.changed.large_communities |= state.large_communities.len() != before;
            Ok(())
        },
    );
}