//! 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::default_policy;
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::errors::{BgpError, CeaseSubcode};
use bgp_packet::message::Message;
use bgp_packet::open::{BgpRole, Capability};
use bgp_primitives::constants::AfiSafi;
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::LocalRoute;
use bgp_server::config::PeerConfig;
use bgp_server::family::FamilySet;
use bgp_server::policy::{ExportLimits, PeerRole, PrefixSet};
use eyre::{Result, bail};
use tokio::net::TcpListener;

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

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

/// What a test holds on to: the server, and the listener its dialling
/// goes to.
struct Setup {
    server: TestServer,
    _listener: TcpListener,
}

/// 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<Setup> {
    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
    let endpoint = listener.local_addr()?;
    let originate = originated
        .iter()
        .map(|p| LocalRoute::new(prefix(p), Default::default()))
        .collect();
    let server = TestServer::start_with(LOCAL_ASN, |c| {
        c.accept_unconfigured_peers = false;
        c.originate = originate;
        c.peers = vec![PeerConfig {
            graceful_restart: Default::default(),
            groups: Vec::new(),
            export_limits: Some(Arc::new(limits)),
            reflector_client: false,
            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,
            export: Default::default(),
        }];
    })
    .await?;
    Ok(Setup {
        server,
        _listener: listener,
    })
}

/// A session with the server, as its peer.
async fn established(server: &TestServer) -> Result<TestPeer> {
    let mut peer = TestPeer::connect(server.addr).await?;
    let capabilities = vec![
        Capability::FourByteAsn { asn: PEER_ASN },
        Capability::multiprotocol(AfiSafi::IPV4_UNICAST),
        Capability::multiprotocol(AfiSafi::IPV6_UNICAST),
        Capability::BgpRole {
            role: BgpRole::Peer,
        },
    ];
    peer.handshake_with(PEER_ASN, Ipv4Addr::new(10, 0, 0, 2), 90, capabilities)
        .await?;
    Ok(peer)
}

/// 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(Ok(message)) = tokio::time::timeout(Duration::from_millis(500), peer.recv()).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 Setup { 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 Setup { 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.is_closed().await);

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