//! Passing routes on: what one peer announced reaches the others, as far
//! as the rules for that let it.

use std::collections::BTreeSet;
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::time::Duration;

use bgp_integration_tests::announcement::{announcement, withdrawal};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_integration_tests::{TIMEOUT, prefix};
use bgp_packet::message::{
    CommunitiesPathAttribute, Message, NextHopPathAttribute, PathAttribute, RouteRefreshMessage,
    UpdateMessage,
};
use bgp_primitives::constants::AfiSafi;
use bgp_rib::{LocalRoute, LocalRouteAttrs};
use bgp_server::marks::{Control, Marks};
use bgp_server::policy::ExportPolicy;
use eyre::Result;

const LOCAL_ASN: u32 = 64601;

/// A server that passes on to every peer what it learned from the others.
async fn transit(customize: impl FnOnce(&mut bgp_server::config::Config)) -> Result<TestServer> {
    TestServer::start_with(LOCAL_ASN, |config| {
        config.export = Arc::new(ExportPolicy {
            learned_routes: true,
            all_paths: false,
            ..ExportPolicy::default()
        });
        customize(config);
    })
    .await
}

/// The peer with AS `64700 + n`, its session up.
async fn peer(server: &TestServer, n: u8) -> Result<TestPeer> {
    server
        .established(64700 + u32::from(n), Ipv4Addr::new(10, 0, 0, n))
        .await
}

async fn announce(
    peer: &mut TestPeer,
    prefixes: &[&str],
    path: &[u32],
    communities: &[(u16, u16)],
) -> Result<()> {
    peer.announce(announcement(prefixes, path).with_communities(communities))
        .await
}

async fn withdraw(peer: &mut TestPeer, prefixes: &[&str]) -> Result<()> {
    peer.announce(withdrawal(prefixes.iter().map(|p| prefix(p).into())))
        .await
}

/// What a peer has been told so far, as the prefixes it would now have
/// and the path each came with.
#[derive(Default, Debug)]
struct Told {
    routes: std::collections::BTreeMap<String, Vec<u32>>,
    updates: usize,
}

impl Told {
    fn take(&mut self, update: &UpdateMessage) {
        self.updates += 1;
        for route in &update.withdrawn_routes {
            self.routes.remove(&route.prefix.to_string());
        }
        let path: Vec<u32> = update
            .path_attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::AsPath(path) => Some(
                    path.segments
                        .iter()
                        .flat_map(|segment| segment.path.clone())
                        .collect(),
                ),
                _ => None,
            })
            .unwrap_or_default();
        for route in &update.announced_routes {
            self.routes.insert(route.prefix.to_string(), path.clone());
        }
    }

    fn prefixes(&self) -> BTreeSet<&str> {
        self.routes.keys().map(String::as_str).collect()
    }
}

/// Reads what the peer is sent until it has exactly `expected`.
async fn told_becomes(peer: &mut TestPeer, told: &mut Told, expected: &[&str]) -> Result<()> {
    let expected: BTreeSet<&str> = expected.iter().copied().collect();
    let reading = async {
        while told.prefixes() != expected {
            let update = peer.recv_update().await?;
            told.take(&update);
        }
        Ok::<(), eyre::Report>(())
    };
    match tokio::time::timeout(TIMEOUT, reading).await {
        Ok(read) => read,
        Err(_) => Err(eyre::eyre!(
            "the peer has {:?}; expected {expected:?}",
            told.prefixes()
        )),
    }
}

/// Nothing more is sent to the peer for a while.
async fn told_nothing(peer: &mut TestPeer) {
    let more = tokio::time::timeout(Duration::from_millis(300), peer.recv_update()).await;
    assert!(more.is_err(), "the peer was sent {more:?}");
}

/// What one peer announces the others are told, with our AS on the path
/// and ourselves as next hop; and what it takes back, they are told too.
#[tokio::test]
async fn test_routes_from_one_peer_reach_the_others() -> Result<()> {
    let server = transit(|_| {}).await?;
    let mut one = peer(&server, 1).await?;
    let mut two = peer(&server, 2).await?;

    announce(
        &mut one,
        &["10.1.0.0/16", "10.2.0.0/16", "10.3.0.0/16"],
        &[64701, 65000],
        &[],
    )
    .await?;
    let mut told = Told::default();
    told_becomes(
        &mut two,
        &mut told,
        &["10.1.0.0/16", "10.2.0.0/16", "10.3.0.0/16"],
    )
    .await?;
    assert_eq!(told.routes["10.1.0.0/16"], [LOCAL_ASN, 64701, 65000]);
    assert_eq!(
        told.updates, 1,
        "routes that go out alike go in one message"
    );
    // Not back to where they came from.
    told_nothing(&mut one).await;

    // The next hop is ourselves: the address of our end of the session.
    announce(&mut one, &["10.4.0.0/16"], &[64701], &[]).await?;
    let update = two.recv_update().await?;
    assert!(
        update
            .path_attributes
            .contains(&PathAttribute::NextHop(NextHopPathAttribute(
                Ipv4Addr::LOCALHOST
            ))),
        "{update:?}"
    );
    told.take(&update);

    // One is taken back; then the peer goes, and its routes with it.
    withdraw(&mut one, &["10.2.0.0/16"]).await?;
    told_becomes(
        &mut two,
        &mut told,
        &["10.1.0.0/16", "10.3.0.0/16", "10.4.0.0/16"],
    )
    .await?;
    drop(one);
    told_becomes(&mut two, &mut told, &[]).await?;
    Ok(())
}

/// A peer is told of the best route, and of the next best when that one
/// goes; and a peer that comes later is told what there is.
#[tokio::test]
async fn test_a_peer_is_kept_at_the_best_route() -> Result<()> {
    let server = transit(|_| {}).await?;
    let mut near = peer(&server, 1).await?;
    let mut far = peer(&server, 2).await?;
    let mut watcher = peer(&server, 3).await?;
    let mut told = Told::default();

    announce(&mut far, &["10.1.0.0/16"], &[64702, 65000, 65001], &[]).await?;
    told_becomes(&mut watcher, &mut told, &["10.1.0.0/16"]).await?;
    assert_eq!(told.routes["10.1.0.0/16"], [LOCAL_ASN, 64702, 65000, 65001]);

    // A better route for it comes: that is the one told now.
    announce(&mut near, &["10.1.0.0/16"], &[64701, 65001], &[]).await?;
    let update = watcher.recv_update().await?;
    told.take(&update);
    assert_eq!(told.routes["10.1.0.0/16"], [LOCAL_ASN, 64701, 65001]);
    assert!(
        update.withdrawn_routes.is_empty(),
        "replaced, not withdrawn"
    );

    // And goes again: back to the other, with no gap.
    withdraw(&mut near, &["10.1.0.0/16"]).await?;
    let update = watcher.recv_update().await?;
    told.take(&update);
    assert_eq!(told.routes["10.1.0.0/16"], [LOCAL_ASN, 64702, 65000, 65001]);

    // The other peer was told of the better route while it was the best,
    // and then that it is gone: the one left is its own, which it is not
    // told of.
    let mut told_far = Told::default();
    told_becomes(&mut far, &mut told_far, &["10.1.0.0/16"]).await?;
    let gone = far.recv_update().await?;
    assert_eq!(gone.withdrawn_routes, vec![prefix("10.1.0.0/16")]);
    told_nothing(&mut far).await;

    // A peer that comes now is told what there is; and again if it asks.
    let mut late = peer(&server, 4).await?;
    let mut late_told = Told::default();
    told_becomes(&mut late, &mut late_told, &["10.1.0.0/16"]).await?;
    late.send(&Message::RouteRefresh(RouteRefreshMessage::request(
        AfiSafi::IPV4_UNICAST,
    )))
    .await?;
    let again = late.recv_update().await?;
    assert_eq!(again.announced_routes, vec![prefix("10.1.0.0/16")]);
    Ok(())
}

/// A route is not sent to a peer whose AS is on its path already, nor
/// anywhere if it says NO_EXPORT; and what we originate goes to everyone
/// beside what was learned.
#[tokio::test]
async fn test_routes_stop_where_they_must() -> Result<()> {
    let server = transit(|config| {
        config.originate = vec![LocalRoute::new(
            prefix("10.99.0.0/16"),
            LocalRouteAttrs::default(),
        )];
    })
    .await?;
    let mut one = peer(&server, 1).await?;
    let mut two = peer(&server, 2).await?;
    let mut told = Told::default();
    told_becomes(&mut two, &mut told, &["10.99.0.0/16"]).await?;
    assert_eq!(told.routes["10.99.0.0/16"], [LOCAL_ASN]);

    // Through the second peer's AS already; and not for other networks.
    announce(&mut one, &["10.1.0.0/16"], &[64701, 64702, 65000], &[]).await?;
    announce(&mut one, &["10.2.0.0/16"], &[64701], &[(0xFFFF, 0xFF01)]).await?;
    announce(&mut one, &["10.3.0.0/16"], &[64701], &[]).await?;
    told_becomes(&mut two, &mut told, &["10.99.0.0/16", "10.3.0.0/16"]).await?;
    told_nothing(&mut two).await;
    assert_eq!(server.rib.len(), 3, "all three were taken");
    Ok(())
}

/// Communities on a route say which peers it is passed on to, and are
/// gone when it is; and every route to a peer gets the marks of `export`.
#[tokio::test]
async fn test_communities_say_who_is_told() -> Result<()> {
    let server = transit(|config| {
        config.control = Some(Control {
            local_asn: LOCAL_ASN,
        });
        config.export = Arc::new(ExportPolicy {
            learned_routes: true,
            all_paths: false,
            marks: Marks::parse(&[], &["64601:7".to_string()], &[]).unwrap(),
            prepend: 1,
        });
        config.originate = vec![LocalRoute::new(
            prefix("10.99.0.0/16"),
            LocalRouteAttrs {
                // Our own route, and not for the second peer.
                communities: vec![(0, 64702)],
                ..LocalRouteAttrs::default()
            },
        )];
    })
    .await?;
    let mut one = peer(&server, 1).await?;
    let mut two = peer(&server, 2).await?;
    let mut three = peer(&server, 3).await?;
    let (mut told_two, mut told_three) = (Told::default(), Told::default());

    // To everyone but the second; to the second only; to nobody; and
    // with nothing said.
    announce(
        &mut one,
        &["10.1.0.0/16"],
        &[64701],
        &[(0, 64702), (3356, 3)],
    )
    .await?;
    announce(&mut one, &["10.2.0.0/16"], &[64701], &[(64601, 64702)]).await?;
    announce(&mut one, &["10.3.0.0/16"], &[64701], &[(0, 64601)]).await?;
    announce(&mut one, &["10.4.0.0/16"], &[64701], &[]).await?;

    told_becomes(&mut two, &mut told_two, &["10.2.0.0/16", "10.4.0.0/16"]).await?;
    told_becomes(
        &mut three,
        &mut told_three,
        &["10.99.0.0/16", "10.1.0.0/16", "10.4.0.0/16"],
    )
    .await?;
    told_nothing(&mut two).await;
    told_nothing(&mut three).await;
    // Our AS once for passing it on, and once more by `export`.
    assert_eq!(
        told_three.routes["10.1.0.0/16"],
        [LOCAL_ASN, LOCAL_ASN, 64701]
    );
    assert_eq!(told_three.routes["10.99.0.0/16"], [LOCAL_ASN, LOCAL_ASN]);

    // What was for us to act on is gone, and the router's mark is on.
    announce(
        &mut one,
        &["10.5.0.0/16"],
        &[64701],
        &[(0, 64702), (3356, 3)],
    )
    .await?;
    let update = three.recv_update().await?;
    assert!(
        update
            .path_attributes
            .contains(&PathAttribute::Communities(CommunitiesPathAttribute(vec![
                (3356, 3),
                (64601, 7)
            ]))),
        "{update:?}"
    );
    Ok(())
}

/// Without `learned_routes` a peer is told what we originate and nothing
/// else, as before there was such a thing.
#[tokio::test]
async fn test_nothing_learned_is_passed_on_unless_asked_for() -> Result<()> {
    let server = TestServer::start_with(LOCAL_ASN, |config| {
        config.originate = vec![LocalRoute::new(
            prefix("10.99.0.0/16"),
            LocalRouteAttrs::default(),
        )];
    })
    .await?;
    let mut one = peer(&server, 1).await?;
    let mut two = peer(&server, 2).await?;
    let mut told = Told::default();
    told_becomes(&mut two, &mut told, &["10.99.0.0/16"]).await?;
    announce(&mut one, &["10.1.0.0/16"], &[64701], &[]).await?;
    server.wait_for_routes(1).await?;
    told_nothing(&mut two).await;
    Ok(())
}