//! End-to-end tests for locally originated routes: announcing via the
//! `LocalRib` handle must produce UPDATEs on every established session.

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

use bgp_integration_tests::announcement::announcement;
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::{RouteChange, TestServer};
use bgp_integration_tests::{TIMEOUT, prefix};
use bgp_packet::message::{
    AsPathAttribute, CommunitiesPathAttribute, LargeCommunitiesPathAttribute, LargeCommunity,
    LocalPrefPathAttribute, NextHopPathAttribute, OriginPathAttribute, PathAttribute,
    UpdateMessage,
};
use bgp_rib::{LocalRoute, LocalRouteAttrs};
use bgp_server::session::SessionEvent;
use eyre::Result;

const LOCAL_ASN: u32 = 65000;
const PEER_ASN: u32 = 65001;

fn route(p: &str, attrs: LocalRouteAttrs) -> LocalRoute {
    LocalRoute::new(prefix(p), attrs)
}

async fn established_peer(server: &TestServer) -> Result<TestPeer> {
    established_peer_with_id(server, Ipv4Addr::new(10, 0, 0, 1)).await
}

async fn established_peer_with_id(server: &TestServer, router_id: Ipv4Addr) -> Result<TestPeer> {
    server.established(PEER_ASN, router_id).await
}

async fn next_update(peer: &mut TestPeer) -> Result<UpdateMessage> {
    peer.recv_update().await
}

// ---------------------------------------------------------------------------
// Session registry
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_registry_tracks_session_lifecycle() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut events = server.registry.subscribe();
    assert!(server.registry.list().is_empty());

    let peer = established_peer(&server).await?;

    let up = tokio::time::timeout(TIMEOUT, events.recv()).await??;
    let SessionEvent::Up(handle) = up else {
        panic!("expected Up, got {up:?}")
    };
    assert_eq!(handle.peer.asn, PEER_ASN);
    assert_eq!(handle.local_addr, server.addr);
    assert!(handle.four_octet_asn);
    assert_eq!(server.registry.list().len(), 1);

    drop(peer);

    let down = tokio::time::timeout(TIMEOUT, events.recv()).await??;
    assert!(
        matches!(down, SessionEvent::Down(p) if p == handle.peer),
        "got {down:?}"
    );
    assert!(server.registry.list().is_empty());
    Ok(())
}

// ---------------------------------------------------------------------------
// Announce / withdraw
// ---------------------------------------------------------------------------

/// A route announced while the session is up is sent immediately, with the
/// attributes the announcer chose and our address as the next hop.
#[tokio::test]
async fn test_announce_after_session_up() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = established_peer(&server).await?;
    server.wait_for_sessions(1).await?;

    server
        .local_rib
        .announce(route(
            "10.100.1.0/24",
            LocalRouteAttrs {
                origin: OriginPathAttribute::IGP,
                prepend: 0,
                med: None,
                communities: vec![(65000, 100), (65535, 65281)],
                large_communities: vec![LargeCommunity {
                    global_admin: 65000,
                    data_1: 7,
                    data_2: 8,
                }],
            },
        ))
        .await?;

    let update = next_update(&mut peer).await?;
    assert_eq!(update.announced_routes, vec![prefix("10.100.1.0/24")]);
    assert!(update.withdrawn_routes.is_empty());
    assert_eq!(
        update.path_attributes,
        vec![
            PathAttribute::Origin(OriginPathAttribute::IGP),
            AsPathAttribute::from_asns(vec![LOCAL_ASN]),
            PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::LOCALHOST)),
            PathAttribute::Communities(CommunitiesPathAttribute(vec![
                (65000, 100),
                (65535, 65281)
            ])),
            PathAttribute::LargeCommunities(LargeCommunitiesPathAttribute {
                communities: vec![LargeCommunity {
                    global_admin: 65000,
                    data_1: 7,
                    data_2: 8
                }],
            }),
        ]
    );
    Ok(())
}

/// Routes announced before any session exists are replayed when a session
/// reaches Established.
#[tokio::test]
async fn test_routes_replayed_on_session_up() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    server
        .local_rib
        .announce(route("10.100.1.0/24", LocalRouteAttrs::default()))
        .await?;
    server
        .local_rib
        .announce(route("10.100.2.0/24", LocalRouteAttrs::default()))
        .await?;

    let mut peer = established_peer(&server).await?;

    let mut seen = vec![next_update(&mut peer).await?, next_update(&mut peer).await?]
        .into_iter()
        .flat_map(|u| u.announced_routes)
        .map(|p| p.to_string())
        .collect::<Vec<_>>();
    seen.sort();
    assert_eq!(seen, vec!["10.100.1.0/24", "10.100.2.0/24"]);
    Ok(())
}

/// An announcement is fanned out to every established session.
#[tokio::test]
async fn test_announce_reaches_all_sessions() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut a = established_peer_with_id(&server, Ipv4Addr::new(10, 0, 0, 1)).await?;
    let mut b = established_peer_with_id(&server, Ipv4Addr::new(10, 0, 0, 2)).await?;
    server.wait_for_sessions(2).await?;

    server
        .local_rib
        .announce(route("10.100.1.0/24", LocalRouteAttrs::default()))
        .await?;

    for peer in [&mut a, &mut b] {
        let update = next_update(peer).await?;
        assert_eq!(update.announced_routes, vec![prefix("10.100.1.0/24")]);
    }
    Ok(())
}

#[tokio::test]
async fn test_withdraw_sends_withdrawn_route() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = established_peer(&server).await?;
    server.wait_for_sessions(1).await?;

    server
        .local_rib
        .announce(route("10.100.1.0/24", LocalRouteAttrs::default()))
        .await?;
    let _ = next_update(&mut peer).await?;

    server.local_rib.withdraw(prefix("10.100.1.0/24")).await?;
    let update = next_update(&mut peer).await?;
    assert_eq!(update.withdrawn_routes, vec![prefix("10.100.1.0/24")]);
    assert!(update.announced_routes.is_empty());
    assert!(update.path_attributes.is_empty());
    Ok(())
}

/// Re-announcing with new attributes sends a fresh UPDATE (implicit replace).
#[tokio::test]
async fn test_reannounce_with_prepend_and_med() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = established_peer(&server).await?;
    server.wait_for_sessions(1).await?;

    server
        .local_rib
        .announce(route("10.100.1.0/24", LocalRouteAttrs::default()))
        .await?;
    let first = next_update(&mut peer).await?;
    assert_eq!(
        first.path_attributes[1],
        AsPathAttribute::from_asns(vec![LOCAL_ASN])
    );

    server
        .local_rib
        .announce(route(
            "10.100.1.0/24",
            LocalRouteAttrs {
                prepend: 3,
                med: Some(42),
                ..LocalRouteAttrs::default()
            },
        ))
        .await?;
    let second = next_update(&mut peer).await?;
    assert_eq!(second.announced_routes, vec![prefix("10.100.1.0/24")]);
    assert_eq!(
        second.path_attributes[1],
        AsPathAttribute::from_asns(vec![LOCAL_ASN; 4])
    );
    assert!(
        second
            .path_attributes
            .iter()
            .any(|a| matches!(a, PathAttribute::MultiExitDisc(m) if m.0 == 42))
    );
    Ok(())
}

/// The session stays healthy (keepalives keep flowing) while UPDATEs are sent.
#[tokio::test]
async fn test_keepalives_continue_alongside_updates() -> Result<()> {
    let server = TestServer::start_with_hold_time(LOCAL_ASN, 3).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(PEER_ASN, Ipv4Addr::new(10, 0, 0, 1), 3)
        .await?;
    server.wait_for_sessions(1).await?;

    server
        .local_rib
        .announce(route("10.100.1.0/24", LocalRouteAttrs::default()))
        .await?;

    let mut saw_update = false;
    let mut saw_keepalive = false;
    tokio::time::timeout(TIMEOUT, async {
        while !(saw_update && saw_keepalive) {
            match peer.recv().await? {
                bgp_packet::message::Message::Update(_) => saw_update = true,
                bgp_packet::message::Message::KeepAlive => saw_keepalive = true,
                other => eyre::bail!("unexpected {other:?}"),
            }
        }
        Ok::<_, eyre::Error>(())
    })
    .await??;
    Ok(())
}

/// Routes learned from a peer disappear from the Adj-RIB-In when its
/// session ends, with a Withdrawn event for each.
#[tokio::test]
async fn test_received_routes_removed_when_session_closes() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut rib_events = server.route_changes();
    let mut session_events = server.registry.subscribe();
    let mut peer = established_peer(&server).await?;
    server.wait_for_sessions(1).await?;

    peer.announce(
        announcement(&["203.0.113.0/24", "198.51.100.0/24"], &[PEER_ASN])
            .with_next_hop(Ipv4Addr::new(10, 0, 0, 1)),
    )
    .await?;
    for _ in 0..2 {
        let ev = tokio::time::timeout(TIMEOUT, rib_events.recv()).await??;
        assert!(matches!(ev, RouteChange::Announced { .. }), "{ev:?}");
    }
    assert_eq!(server.received_routes().len(), 2);

    drop(peer);

    let mut withdrawn = vec![];
    for _ in 0..2 {
        match tokio::time::timeout(TIMEOUT, rib_events.recv()).await?? {
            RouteChange::Withdrawn { prefix, .. } => withdrawn.push(prefix.to_string()),
            other => panic!("expected Withdrawn, got {other:?}"),
        }
    }
    withdrawn.sort();
    assert_eq!(withdrawn, vec!["198.51.100.0/24", "203.0.113.0/24"]);
    assert!(server.received_routes().is_empty());

    // Withdrawals precede the Down event.
    let _up = tokio::time::timeout(TIMEOUT, session_events.recv()).await??;
    let down = tokio::time::timeout(TIMEOUT, session_events.recv()).await??;
    assert!(matches!(down, SessionEvent::Down(_)), "{down:?}");
    Ok(())
}

// ---------------------------------------------------------------------------
// Route refresh and shutdown
// ---------------------------------------------------------------------------

/// ROUTE-REFRESH for IPv4 unicast exactly as a router sends it (RFC 2918):
/// 19-byte header plus AFI, reserved, SAFI.
#[rustfmt::skip]
const ROUTE_REFRESH_IPV4: &[u8] = &[
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0x00, 0x17, 0x05,
    0x00, 0x01, 0x00, 0x01,
];

/// A peer whose import policy changed asks for our routes again; it must get
/// every local route re-sent, and the session must stay up.
#[tokio::test]
async fn test_route_refresh_resends_local_routes() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = established_peer(&server).await?;
    server.wait_for_sessions(1).await?;

    server
        .local_rib
        .announce(route("10.100.1.0/24", LocalRouteAttrs::default()))
        .await?;
    server
        .local_rib
        .announce(route("10.100.2.0/24", LocalRouteAttrs::default()))
        .await?;
    let _ = next_update(&mut peer).await?;
    let _ = next_update(&mut peer).await?;

    peer.send_raw(ROUTE_REFRESH_IPV4).await?;

    let mut resent = vec![next_update(&mut peer).await?, next_update(&mut peer).await?]
        .into_iter()
        .flat_map(|u| u.announced_routes)
        .map(|p| p.to_string())
        .collect::<Vec<_>>();
    resent.sort();
    assert_eq!(resent, vec!["10.100.1.0/24", "10.100.2.0/24"]);
    assert_eq!(
        server.registry.list().len(),
        1,
        "session must survive a route refresh"
    );
    Ok(())
}

/// Shutting down tells each peer why (Cease / Administrative Shutdown) and
/// closes the sessions, instead of leaving dead TCP connections behind.
#[tokio::test]
async fn test_shutdown_sends_cease_to_every_peer() -> Result<()> {
    use bgp_packet::errors::{BgpError, CeaseSubcode};

    let server = TestServer::start(LOCAL_ASN).await?;
    let mut a = established_peer_with_id(&server, Ipv4Addr::new(10, 0, 0, 1)).await?;
    let mut b = established_peer_with_id(&server, Ipv4Addr::new(10, 0, 0, 2)).await?;
    server.wait_for_sessions(2).await?;

    let left = server
        .registry
        .leave(bgp_server::restart::Leave::Stop, Duration::from_secs(2))
        .await;
    assert_eq!(left, 0, "sessions still open after shutdown");
    assert!(server.registry.is_shutting_down());

    for peer in [&mut a, &mut b] {
        let notif = peer.expect_notification().await?;
        assert!(
            matches!(notif.error, BgpError::Cease(CeaseSubcode::AdminShutdown)),
            "{notif:?}"
        );
        assert!(peer.closed_within(TIMEOUT).await);
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Routes from the configuration
// ---------------------------------------------------------------------------

/// Routes listed under `originate` are in the table before any peer
/// connects, and each session is sent them when it comes up.
#[tokio::test]
async fn test_configured_routes_are_announced_to_every_new_session() -> Result<()> {
    let attrs = LocalRouteAttrs {
        med: Some(7),
        communities: vec![(65000, 100)],
        ..Default::default()
    };
    let configured = vec![
        route("10.100.1.0/24", attrs),
        route("10.100.2.0/24", LocalRouteAttrs::default()),
    ];
    let server = TestServer::start_with(LOCAL_ASN, |c| c.originate = configured).await?;
    assert_eq!(server.local_rib.snapshot().await.unwrap().len(), 2);

    for router_id in [Ipv4Addr::new(10, 0, 0, 1), Ipv4Addr::new(10, 0, 0, 2)] {
        let mut peer = established_peer_with_id(&server, router_id).await?;
        let mut seen = [next_update(&mut peer).await?, next_update(&mut peer).await?];
        seen.sort_by_key(|u| u.announced_routes[0].to_string());
        assert_eq!(seen[0].announced_routes, vec![prefix("10.100.1.0/24")]);
        assert!(
            seen[0]
                .path_attributes
                .contains(&PathAttribute::Communities(CommunitiesPathAttribute(vec![
                    (65000, 100)
                ])))
        );
        assert_eq!(seen[1].announced_routes, vec![prefix("10.100.2.0/24")]);
    }
    Ok(())
}

/// A peer in our own AS is sent a route we originate with an empty AS path
/// and a LOCAL_PREF, whatever prepending is configured: our AS number goes
/// on the path only as the route leaves the AS. It used to be sent as to
/// an external peer, which an internal one refuses.
#[tokio::test]
async fn test_a_peer_in_our_own_as_gets_an_empty_path_and_a_local_pref() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(LOCAL_ASN, Ipv4Addr::new(10, 0, 0, 1), 90)
        .await?;
    // The peer reads as a router in the server's AS does: one outside it
    // would drop the LOCAL_PREF.
    peer.ctx.ibgp = true;
    server.wait_for_sessions(1).await?;
    assert!(server.registry.list()[0].ibgp);

    server
        .local_rib
        .announce(route(
            "10.100.1.0/24",
            LocalRouteAttrs {
                prepend: 2,
                ..LocalRouteAttrs::default()
            },
        ))
        .await?;
    let update = next_update(&mut peer).await?;
    assert_eq!(update.announced_routes, vec![prefix("10.100.1.0/24")]);
    assert!(
        update
            .path_attributes
            .contains(&PathAttribute::AsPath(AsPathAttribute { segments: vec![] })),
        "{:?}",
        update.path_attributes
    );
    assert!(
        update
            .path_attributes
            .contains(&PathAttribute::LocalPref(LocalPrefPathAttribute(100))),
        "{:?}",
        update.path_attributes
    );
    Ok(())
}