//! 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::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::message::{
    AsPathAttribute, CommunitiesPathAttribute, LargeCommunitiesPathAttribute, LargeCommunity,
    NextHopPathAttribute, OriginPathAttribute, PathAttribute, UpdateMessage,
};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::{LocalRibError, LocalRibPolicy, LocalRoute, LocalRouteAttrs};
use bgp_server::session::SessionEvent;
use eyre::Result;

const LOCAL_ASN: u32 = 65000;
const PEER_ASN: u32 = 65001;
const TIMEOUT: Duration = Duration::from_secs(5);

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

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

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

/// Distinct routers need distinct BGP Identifiers, or the second connection
/// is (correctly) treated as a collision with the first.
async fn established_peer_with_id(server: &TestServer, router_id: Ipv4Addr) -> Result<TestPeer> {
    let mut peer = TestPeer::connect(server.addr).await?;
    peer.handshake(PEER_ASN, router_id, 90).await?;
    Ok(peer)
}

async fn next_update(peer: &mut TestPeer) -> Result<UpdateMessage> {
    tokio::time::timeout(TIMEOUT, peer.recv_update()).await?
}

/// Waits until the registry reports `n` established sessions.
async fn wait_for_sessions(server: &TestServer, n: usize) -> Result<()> {
    tokio::time::timeout(TIMEOUT, async {
        while server.registry.list().len() != n {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    Ok(())
}

// ---------------------------------------------------------------------------
// 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?;
    wait_for_sessions(&server, 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?;
    wait_for_sessions(&server, 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?;
    wait_for_sessions(&server, 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"), Some("alice")).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?;
    wait_for_sessions(&server, 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?;
    wait_for_sessions(&server, 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.adj_rib_in.subscribe();
    let mut session_events = server.registry.subscribe();
    let mut peer = established_peer(&server).await?;
    wait_for_sessions(&server, 1).await?;

    peer.send(&bgp_packet::message::Message::Update(bgp_packet::message::UpdateMessage {
        withdrawn_routes: vec![],
        path_attributes: vec![
            PathAttribute::Origin(OriginPathAttribute::IGP),
            AsPathAttribute::from_asns(vec![PEER_ASN]),
            PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(10, 0, 0, 1))),
        ],
        announced_routes: vec![prefix("203.0.113.0/24"), prefix("198.51.100.0/24")],
    }))
    .await?;
    for _ in 0..2 {
        let ev = tokio::time::timeout(TIMEOUT, rib_events.recv()).await??;
        assert!(matches!(ev, bgp_rib::RibEvent::Announced { .. }), "{ev:?}");
    }
    assert_eq!(server.adj_rib_in.snapshot().await.len(), 2);

    drop(peer);

    let mut withdrawn = vec![];
    for _ in 0..2 {
        match tokio::time::timeout(TIMEOUT, rib_events.recv()).await?? {
            bgp_rib::RibEvent::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.adj_rib_in.snapshot().await.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?;
    wait_for_sessions(&server, 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};
    use bgp_packet::message::Message;

    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?;
    wait_for_sessions(&server, 2).await?;

    let left = server.registry.shutdown(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 = tokio::time::timeout(TIMEOUT, async {
            loop {
                match peer.recv().await? {
                    Message::Notification(n) => return Ok::<_, eyre::Error>(n),
                    Message::KeepAlive => {}
                    other => eyre::bail!("unexpected {other:?}"),
                }
            }
        })
        .await??;
        assert!(matches!(notif.error, BgpError::Cease(CeaseSubcode::AdminShutdown)), "{notif:?}");
        assert!(peer.is_closed().await);
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Policy
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_policy_rejects_prefix_outside_allowlist() -> Result<()> {
    let server = TestServer::start_with_policy(
        LOCAL_ASN,
        LocalRibPolicy { allowed_prefixes: vec![prefix("10.100.0.0/16")], max_routes_per_owner: 2 },
    )
    .await?;
    let mut peer = established_peer(&server).await?;
    wait_for_sessions(&server, 1).await?;

    let err = server
        .local_rib
        .announce(route("192.0.2.0/24", LocalRouteAttrs::default()))
        .await
        .unwrap_err();
    assert_eq!(err, LocalRibError::NotAllowed(prefix("192.0.2.0/24")));

    // An allowed prefix still goes through on the same session.
    server.local_rib.announce(route("10.100.9.0/24", LocalRouteAttrs::default())).await?;
    let update = next_update(&mut peer).await?;
    assert_eq!(update.announced_routes, vec![prefix("10.100.9.0/24")]);
    Ok(())
}