//! The operator's own rules, as a script: what they do to routes on their
//! way in, and to what a peer is told.

use std::collections::BTreeMap;
use std::net::Ipv4Addr;
use std::time::Duration;

use bgp_integration_tests::default_policy;
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::message::{
    AsPathAttribute, CommunitiesPathAttribute, LocalPrefPathAttribute, Message,
    NextHopPathAttribute, OriginPathAttribute, PathAttribute, UpdateMessage,
};
use bgp_policy::Policy;
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::{LocalRoute, LocalRouteAttrs};
use bgp_server::config::PeerConfig;
use bgp_server::family::FamilySet;
use bgp_server::policy::PeerRole;
use eyre::Result;
use tokio::net::TcpListener;

const LOCAL_ASN: u32 = 1299;
const PEER_ASN: u32 = 3333;
const TIMEOUT: Duration = Duration::from_secs(5);

const SCRIPT: &str = r#"
    const BLACKHOLE = "65535:666";
    const LEARNED_FROM = #{ customer: "1299:1000", peer: "1299:2000", provider: "1299:3000" };

    fn from_peer(route, peer) {
        if route.has_community(BLACKHOLE) { return false; }
        // Nobody else's word for where we learned a route.
        route.remove_community("1299:*");
        route.add_community(LEARNED_FROM[peer.role]);
        route.local_pref = if route.has_community("3333:80") { 80 } else { 200 };
        true
    }

    fn by_prefix(route, peer) {
        if route.within("194.0.0.0/8") { return false; }
        if route.prefix == "193.0.4.0/24" { route.add_community("1299:4"); }
        true
    }

    fn broken(route, peer) {
        if route.prefix_len == 22 { return route.no_such_thing; }
        true
    }

    fn to_peer(route, peer) {
        if route.has_community("1299:9") { return false; }      // not for anyone else
        if route.within("193.0.0.0/21") { route.prepend(2); }
        route.add_community("1299:" + peer.asn % 1000);
        route.med = 50;
        true
    }
"#;

fn prefix(s: &str) -> IpPrefix {
    IpPrefix::try_from(s).unwrap()
}

/// A server with one configured peer, whose routes are given to the
/// script's function `import`, and which is told what `export` lets by.
async fn server_with(
    import: Option<&str>,
    export: Option<&str>,
    originate: Vec<LocalRoute>,
) -> Result<(TestServer, TcpListener)> {
    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
    let endpoint = listener.local_addr()?;
    let policy = Policy::compile(SCRIPT)?;
    let import_hook = import.map(|name| policy.hook(name)).transpose()?;
    let export_hook = export.map(|name| policy.hook(name)).transpose()?;
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        c.accept_unconfigured_peers = false;
        c.originate = originate;
        c.policy = Some(policy);
        c.peers = vec![PeerConfig {
            name: "r1".into(),
            endpoint: endpoint.into(),
            remote_asn: PEER_ASN,
            hold_time_secs: 90,
            local_address: None,
            families: FamilySet::supported(),
            add_path: None,
            import: default_policy(),
            role: Some(PeerRole::Peer),
            require_role: false,
            import_hook,
            export_hook,
        }];
    })
    .await?;
    Ok((server, listener))
}

async fn established(server: &TestServer) -> Result<TestPeer> {
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90)
        .await?;
    tokio::time::timeout(TIMEOUT, async {
        while server.registry.list().is_empty() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await?;
    Ok(peer)
}

async fn announce(
    peer: &mut TestPeer,
    prefixes: &[&str],
    communities: &[(u16, u16)],
) -> Result<()> {
    let mut path_attributes = vec![
        PathAttribute::Origin(OriginPathAttribute::IGP),
        AsPathAttribute::from_asns(vec![PEER_ASN]),
        PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(192, 0, 2, 1))),
    ];
    if !communities.is_empty() {
        path_attributes.push(PathAttribute::Communities(CommunitiesPathAttribute(
            communities.to_vec(),
        )));
    }
    peer.send(&Message::Update(UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes,
        announced_routes: prefixes.iter().map(|p| prefix(p).into()).collect(),
    }))
    .await
}

/// Waits until the server has taken `taken` routes of the session's and
/// refused `refused`.
async fn dealt_with(server: &TestServer, taken: u64, refused: u64) -> Result<()> {
    let counters = server.registry.list()[0].counters.clone();
    tokio::time::timeout(TIMEOUT, async {
        while (counters.routes(), counters.rejected()) != (taken, refused) {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .map_err(|_| {
        eyre::eyre!(
            "the server took {} routes and refused {}; expected {taken} and {refused}",
            counters.routes(),
            counters.rejected()
        )
    })
}

/// A route's communities and its LOCAL_PREF.
type Marks = (Vec<(u16, u16)>, Option<u32>);

/// The communities and LOCAL_PREF each prefix is in the table with.
fn table(server: &TestServer) -> BTreeMap<String, Marks> {
    server
        .received_routes()
        .iter()
        .map(|route| {
            let attributes = &route.attributes;
            let communities = attributes
                .iter()
                .find_map(|attribute| match attribute {
                    PathAttribute::Communities(communities) => Some(communities.0.clone()),
                    _ => None,
                })
                .unwrap_or_default();
            let local_pref = attributes.iter().find_map(|attribute| match attribute {
                PathAttribute::LocalPref(LocalPrefPathAttribute(value)) => Some(*value),
                _ => None,
            });
            (route.prefix.to_string(), (communities, local_pref))
        })
        .collect()
}

fn refusals(server: &TestServer) -> Vec<(&'static str, u64)> {
    server.refusals.peers()[0]
        .1
        .by_reason()
        .into_iter()
        .map(|(reason, count)| (reason.as_str(), count))
        .collect()
}

/// Routes are marked with where they were learned and given a
/// preference, by what they carry and who sent them; one the function
/// refuses is not taken.
#[tokio::test]
async fn test_routes_are_marked_preferred_and_refused_on_their_way_in() -> Result<()> {
    let (server, _listener) = server_with(Some("from_peer"), None, vec![]).await?;
    let mut peer = established(&server).await?;

    // One UPDATE of two routes; one of a route with marks of its own, one
    // of them forged; and one the function refuses.
    announce(&mut peer, &["193.0.0.0/21", "193.0.8.0/22"], &[]).await?;
    announce(&mut peer, &["194.0.1.0/24"], &[(3333, 80), (1299, 1000)]).await?;
    announce(&mut peer, &["195.0.1.0/24"], &[(65535, 666)]).await?;
    dealt_with(&server, 3, 1).await?;

    assert_eq!(
        table(&server),
        BTreeMap::from([
            ("193.0.0.0/21".to_string(), (vec![(1299, 2000)], Some(200))),
            ("193.0.8.0/22".to_string(), (vec![(1299, 2000)], Some(200))),
            (
                "194.0.1.0/24".to_string(),
                (vec![(3333, 80), (1299, 2000)], Some(80))
            ),
        ])
    );
    assert_eq!(refusals(&server), [("refused by policy function", 1)]);

    // A route the function took is replaced by one it refuses: gone.
    announce(&mut peer, &["193.0.8.0/22"], &[(65535, 666)]).await?;
    dealt_with(&server, 2, 2).await?;
    assert!(!table(&server).contains_key("193.0.8.0/22"));
    Ok(())
}

/// A function that asks for the prefix is asked about each route of an
/// UPDATE, and may answer differently for each.
#[tokio::test]
async fn test_a_function_may_treat_the_routes_of_one_update_differently() -> Result<()> {
    let (server, _listener) = server_with(Some("by_prefix"), None, vec![]).await?;
    let mut peer = established(&server).await?;
    announce(
        &mut peer,
        &[
            "193.0.0.0/21",
            "194.0.1.0/24",
            "193.0.4.0/24",
            "194.0.2.0/24",
            "195.0.0.0/24",
        ],
        &[(3333, 1)],
    )
    .await?;
    dealt_with(&server, 3, 2).await?;
    assert_eq!(
        table(&server),
        BTreeMap::from([
            ("193.0.0.0/21".to_string(), (vec![(3333, 1)], None)),
            (
                "193.0.4.0/24".to_string(),
                (vec![(3333, 1), (1299, 4)], None)
            ),
            ("195.0.0.0/24".to_string(), (vec![(3333, 1)], None)),
        ])
    );
    assert_eq!(
        server.rib.attrs().len(),
        2,
        "two routes share their attributes"
    );
    Ok(())
}

/// A function that fails on a route has refused it, and no other.
#[tokio::test]
async fn test_a_function_that_fails_refuses_the_route() -> Result<()> {
    let (server, _listener) = server_with(Some("broken"), None, vec![]).await?;
    let mut peer = established(&server).await?;
    announce(
        &mut peer,
        &["193.0.0.0/21", "193.0.8.0/22", "195.0.0.0/24"],
        &[],
    )
    .await?;
    dealt_with(&server, 2, 1).await?;
    assert_eq!(refusals(&server), [("policy function failed", 1)]);
    assert!(!table(&server).contains_key("193.0.8.0/22"));
    Ok(())
}

/// What a peer is told is what the function for it lets by, as the
/// function left it.
#[tokio::test]
async fn test_what_a_peer_is_told_goes_through_its_function() -> Result<()> {
    let route = |prefix_text: &str, communities: Vec<(u16, u16)>| {
        LocalRoute::new(
            prefix(prefix_text),
            LocalRouteAttrs {
                communities,
                ..LocalRouteAttrs::default()
            },
        )
    };
    let originate = vec![
        route("193.0.0.0/21", vec![(1299, 1)]),
        route("194.0.0.0/22", vec![]),
        route("195.0.0.0/22", vec![(1299, 9)]),
    ];
    let (server, _listener) = server_with(None, Some("to_peer"), originate).await?;
    let mut peer = established(&server).await?;

    let mut told = BTreeMap::new();
    for _ in 0..2 {
        let update = tokio::time::timeout(TIMEOUT, peer.recv_update()).await??;
        let path = update
            .path_attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::AsPath(path) => Some(path.segments[0].path.clone()),
                _ => None,
            })
            .unwrap();
        let communities = update
            .path_attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::Communities(communities) => Some(communities.0.clone()),
                _ => None,
            })
            .unwrap_or_default();
        let med = update
            .path_attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::MultiExitDisc(med) => Some(med.0),
                _ => None,
            });
        // In the order of their type codes, as any UPDATE of ours.
        let codes: Vec<u8> = update
            .path_attributes
            .iter()
            .map(PathAttribute::type_code)
            .collect();
        assert!(codes.is_sorted(), "{codes:?}");
        // The peer is one: the route is still marked as only for customers.
        assert!(codes.contains(&35), "{codes:?}");
        told.insert(
            update.announced_routes[0].prefix.to_string(),
            (path, communities, med),
        );
    }
    assert_eq!(
        told,
        BTreeMap::from([
            (
                "193.0.0.0/21".to_string(),
                (vec![LOCAL_ASN; 3], vec![(1299, 1), (1299, 333)], Some(50))
            ),
            (
                "194.0.0.0/22".to_string(),
                (vec![LOCAL_ASN], vec![(1299, 333)], Some(50))
            ),
        ])
    );
    // And nothing of the third.
    assert!(
        tokio::time::timeout(Duration::from_millis(300), peer.recv_update())
            .await
            .is_err()
    );
    Ok(())
}