//! Export limits, end to end: what a peer above or beside us is told is
//! held to the prefixes they list and to as many routes as they allow,
//! whatever else would send it more.

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

use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::{ServerWithPeer, TestServer};
use bgp_integration_tests::{TIMEOUT, prefix};
use bgp_packet::errors::{BgpError, CeaseSubcode};
use bgp_packet::message::Message;
use bgp_packet::open::{BgpRole, Capability};
use bgp_primitives::constants::AfiSafi;
use bgp_rib::LocalRoute;
use bgp_server::policy::{ExportLimits, PeerRole, PrefixSet};
use eyre::{Result, bail};

const LOCAL_ASN: u32 = 1299;
const PEER_ASN: u32 = 3333;
const PEER_ID: Ipv4Addr = Ipv4Addr::new(10, 0, 0, 2);

/// A server that originates `originated`, with one configured peer at the
/// address the test connects from, which is a peer of ours held to
/// `limits`.
async fn server_with(originated: &[&str], limits: ExportLimits) -> Result<ServerWithPeer> {
    let originate = originated
        .iter()
        .map(|p| LocalRoute::new(prefix(p), Default::default()))
        .collect();
    TestServer::with_peer(LOCAL_ASN, PEER_ASN, |c, peer| {
        c.originate = originate;
        peer.export_limits = Some(Arc::new(limits));
        peer.role = Some(PeerRole::Peer);
    })
    .await
}

/// A session with the server, as its peer.
async fn established(server: &TestServer) -> Result<TestPeer> {
    let capabilities = vec![
        Capability::FourByteAsn { asn: PEER_ASN },
        Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
        Capability::multiprotocol(AfiSafi::IPV6_UNICAST),
        Capability::BgpRole {
            role: BgpRole::Peer,
        },
    ];
    server
        .established_with(PEER_ASN, PEER_ID, capabilities)
        .await
}

/// What the peer has been told once it has been quiet for a while: the
/// prefixes it holds, and the NOTIFICATION that ended the session, if one
/// did.
async fn told(peer: &mut TestPeer) -> (BTreeSet<String>, Option<(BgpError, Vec<u8>)>) {
    let mut routes = BTreeSet::new();
    while let Ok(message) = peer.recv_within(Duration::from_millis(500)).await {
        match message {
            Message::Update(update) => {
                for route in &update.withdrawn_routes {
                    routes.remove(&route.prefix.to_string());
                }
                for route in &update.announced_routes {
                    routes.insert(route.prefix.to_string());
                }
            }
            Message::Notification(notification) => {
                return (routes, Some((notification.error, notification.data)));
            }
            _ => {}
        }
    }
    (routes, None)
}

fn allowing(prefixes: &[&str]) -> Option<PrefixSet> {
    let mut set = PrefixSet::default();
    for p in prefixes {
        let p = prefix(p);
        set.allow(p, p.length());
    }
    Some(set)
}

/// A route of ours whose prefix is not in the peer's export limits is not
/// announced to it; the one that is, is.
#[tokio::test]
async fn test_a_route_outside_the_export_limits_is_not_announced() -> Result<()> {
    let limits = ExportLimits {
        prefixes: allowing(&["193.0.0.0/21"]),
        max_prefixes: None,
    };
    let ServerWithPeer { server, .. } =
        server_with(&["193.0.0.0/21", "198.51.100.0/24"], limits).await?;
    let mut peer = established(&server).await?;
    let (routes, ended) = told(&mut peer).await;
    assert_eq!(routes, BTreeSet::from(["193.0.0.0/21".to_string()]));
    assert_eq!(ended, None);
    Ok(())
}

/// One route more than the export limit allows is not announced: the
/// session ends with an Administrative Shutdown that says why (RFC 9003),
/// and the peer is kept down, so that it is told nothing rather than too
/// much.
#[tokio::test]
async fn test_going_over_the_export_limit_ends_the_session_and_keeps_the_peer_down() -> Result<()> {
    let limits = ExportLimits {
        prefixes: None,
        max_prefixes: Some(2),
    };
    let ServerWithPeer { server, .. } =
        server_with(&["193.0.0.0/21", "193.0.8.0/21"], limits).await?;
    let mut peer = established(&server).await?;
    let (routes, ended) = told(&mut peer).await;
    assert_eq!(routes.len(), 2);
    assert_eq!(ended, None);

    server
        .local_rib
        .announce(LocalRoute::new(prefix("193.0.16.0/21"), Default::default()))
        .await?;
    let (more, ended) = told(&mut peer).await;
    assert_eq!(more, BTreeSet::new(), "nothing more is announced");
    let Some((error, data)) = ended else {
        bail!("the session was not ended");
    };
    assert_eq!(error, BgpError::Cease(CeaseSubcode::AdminShutdown));
    let text = String::from_utf8(data[1..].to_vec())?;
    assert!(
        text.starts_with("export limit of 2 routes reached"),
        "{text}"
    );
    assert!(peer.closed_within(TIMEOUT).await);

    // The peer is kept down: a new connection is closed unanswered.
    let mut again = TestPeer::connect(server.addr).await?;
    assert!(
        again.closed_unanswered().await,
        "a held-down peer gets no session"
    );
    Ok(())
}