//! Connection collision detection (RFC 4271 §6.8) against a live server.
//!
//! The test server's BGP Identifier is 127.0.0.1. A test peer picks an
//! identifier above or below that to steer which connection must survive.

use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;

use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::errors::{BgpError, CeaseSubcode, OpenMsgSubcode};
use bgp_packet::message::{
    AsPathAttribute, Message, NextHopPathAttribute, NotificationMessage, OriginPathAttribute,
    PathAttribute, UpdateMessage,
};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::{LocalRoute, LocalRouteAttrs};
use bgp_server::collision::Direction;
use bgp_server::config::PeerConfig;
use bgp_server::family::FamilySet;
use bgp_server::policy::ImportPolicy;
use eyre::Result;
use tokio::net::TcpListener;

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

/// Above the server's 127.0.0.1.
const HIGH_ID: Ipv4Addr = Ipv4Addr::new(192, 0, 2, 1);
/// Below the server's 127.0.0.1.
const LOW_ID: Ipv4Addr = Ipv4Addr::new(10, 0, 0, 1);

async fn server_dialing(listener: &TcpListener) -> Result<TestServer> {
    let endpoint: SocketAddr = listener.local_addr()?;
    TestServer::start_with(LOCAL_ASN, |c| {
        c.connect_retry = Duration::from_millis(200);
        c.peers = vec![PeerConfig {
            graceful_restart: Default::default(),
            groups: Vec::new(),
            export_limits: None,
            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: ImportPolicy::permissive(),
            role: None,
            require_role: false,
            export: Default::default(),
        }];
    })
    .await
}

/// Exchanges OPENs only, leaving the server's side in OpenConfirm (it has
/// sent its KEEPALIVE and is waiting for ours).
async fn open_only(peer: &mut TestPeer, router_id: Ipv4Addr) -> Result<()> {
    peer.send_open(PEER_ASN, router_id, 90).await?;
    peer.recv_open().await?;
    Ok(())
}

/// Finishes the handshake begun by [`open_only`].
async fn confirm(peer: &mut TestPeer) -> Result<()> {
    match tokio::time::timeout(TIMEOUT, peer.recv()).await?? {
        Message::KeepAlive => {}
        other => eyre::bail!("expected KEEPALIVE, got {other:?}"),
    }
    peer.send(&Message::KeepAlive).await
}

/// Reads until a NOTIFICATION arrives, skipping KEEPALIVEs.
async fn expect_notification(peer: &mut TestPeer) -> Result<NotificationMessage> {
    tokio::time::timeout(TIMEOUT, async {
        loop {
            match peer.recv().await? {
                Message::Notification(n) => return Ok(n),
                Message::KeepAlive => {}
                other => eyre::bail!("expected NOTIFICATION, got {other:?}"),
            }
        }
    })
    .await?
}

fn is_collision_cease(n: &NotificationMessage) -> bool {
    matches!(
        n.error,
        BgpError::Cease(CeaseSubcode::ConnCollisionResolution)
    )
}

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(())
}

/// Peer has the higher identifier, so the connection the *peer* initiated
/// must survive: the server closes its own outbound connection.
#[tokio::test]
async fn test_simultaneous_open_peer_higher_keeps_inbound() -> Result<()> {
    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
    let server = server_dialing(&listener).await?;

    // Server's outbound connection reaches OpenConfirm.
    let (stream, _) = tokio::time::timeout(TIMEOUT, listener.accept()).await??;
    let mut dialled = TestPeer::from_stream(stream);
    open_only(&mut dialled, HIGH_ID).await?;

    // The peer now also dials the server.
    let mut inbound = TestPeer::connect(server.addr).await?;
    open_only(&mut inbound, HIGH_ID).await?;

    let notif = expect_notification(&mut dialled).await?;
    assert!(is_collision_cease(&notif), "{notif:?}");
    assert!(dialled.is_closed().await);

    confirm(&mut inbound).await?;
    wait_for_sessions(&server, 1).await?;
    let session = server.registry.list().remove(0);
    assert_eq!(session.direction, Direction::Inbound);
    assert_eq!(session.name.as_deref(), Some("r1"));

    // The survivor is fully functional.
    server
        .local_rib
        .announce(LocalRoute::new(
            IpPrefix::try_from("10.100.1.0/24")?,
            LocalRouteAttrs::default(),
        ))
        .await?;
    let update = tokio::time::timeout(TIMEOUT, inbound.recv_update()).await??;
    assert_eq!(update.announced_routes[0].to_string(), "10.100.1.0/24");

    // The dialer notices the surviving session and does not open another.
    let extra = tokio::time::timeout(Duration::from_millis(700), listener.accept()).await;
    assert!(
        extra.is_err(),
        "dialer opened a second connection despite an established session"
    );
    Ok(())
}

/// Peer has the lower identifier, so the connection the *server* initiated
/// must survive: the peer's inbound connection is refused.
#[tokio::test]
async fn test_simultaneous_open_peer_lower_keeps_outbound() -> Result<()> {
    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
    let server = server_dialing(&listener).await?;

    let (stream, _) = tokio::time::timeout(TIMEOUT, listener.accept()).await??;
    let mut dialled = TestPeer::from_stream(stream);
    open_only(&mut dialled, LOW_ID).await?;

    let mut inbound = TestPeer::connect(server.addr).await?;
    open_only(&mut inbound, LOW_ID).await?;

    let notif = expect_notification(&mut inbound).await?;
    assert!(is_collision_cease(&notif), "{notif:?}");
    assert!(inbound.is_closed().await);

    confirm(&mut dialled).await?;
    wait_for_sessions(&server, 1).await?;
    assert_eq!(server.registry.list()[0].direction, Direction::Outbound);
    Ok(())
}

/// The order in which the two OPENs are processed must not change the
/// outcome: here the inbound connection reaches OpenConfirm first.
#[tokio::test]
async fn test_outcome_is_independent_of_arrival_order() -> Result<()> {
    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?;
    let server = server_dialing(&listener).await?;

    let (stream, _) = tokio::time::timeout(TIMEOUT, listener.accept()).await??;
    let mut dialled = TestPeer::from_stream(stream);

    // Inbound first this time; peer is lower so outbound must still win.
    let mut inbound = TestPeer::connect(server.addr).await?;
    open_only(&mut inbound, LOW_ID).await?;
    open_only(&mut dialled, LOW_ID).await?;

    let notif = expect_notification(&mut inbound).await?;
    assert!(is_collision_cease(&notif), "{notif:?}");

    confirm(&mut dialled).await?;
    wait_for_sessions(&server, 1).await?;
    assert_eq!(server.registry.list()[0].direction, Direction::Outbound);
    Ok(())
}

/// RFC default: a collision with an Established session closes the new
/// connection, whatever the identifiers say.
#[tokio::test]
async fn test_established_session_wins_by_default() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;

    let mut first = TestPeer::connect(server.addr).await?;
    first.handshake(PEER_ASN, HIGH_ID, 90).await?;
    wait_for_sessions(&server, 1).await?;

    let mut second = TestPeer::connect(server.addr).await?;
    open_only(&mut second, HIGH_ID).await?;
    let notif = expect_notification(&mut second).await?;
    assert!(is_collision_cease(&notif), "{notif:?}");

    // The original is untouched.
    assert_eq!(server.registry.list().len(), 1);
    server
        .local_rib
        .announce(LocalRoute::new(
            IpPrefix::try_from("10.100.1.0/24")?,
            LocalRouteAttrs::default(),
        ))
        .await?;
    let update = tokio::time::timeout(TIMEOUT, first.recv_update()).await??;
    assert_eq!(update.announced_routes[0].to_string(), "10.100.1.0/24");
    Ok(())
}

/// With collision_detect_established_state, a peer that restarts and dials
/// again replaces the stale session at once, and the routes learned over the
/// old session are flushed.
#[tokio::test]
async fn test_restarted_peer_replaces_stale_session_when_enabled() -> Result<()> {
    let server =
        TestServer::start_with(LOCAL_ASN, |c| c.collision_detect_established_state = true).await?;

    let mut stale = TestPeer::connect(server.addr).await?;
    stale.handshake(PEER_ASN, HIGH_ID, 90).await?;
    wait_for_sessions(&server, 1).await?;
    stale
        .send(&Message::Update(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, 9))),
            ],
            announced_routes: vec![IpPrefix::try_from("203.0.113.0/24")?.into()],
        }))
        .await?;
    tokio::time::timeout(TIMEOUT, async {
        while server.received_routes().is_empty() {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    let stale_addr = server.registry.list()[0].peer.addr;

    // "Restart": a fresh connection from the same router.
    let mut fresh = TestPeer::connect(server.addr).await?;
    fresh.handshake(PEER_ASN, HIGH_ID, 90).await?;

    let notif = expect_notification(&mut stale).await?;
    assert!(is_collision_cease(&notif), "{notif:?}");

    tokio::time::timeout(TIMEOUT, async {
        loop {
            let sessions = server.registry.list();
            if sessions.len() == 1 && sessions[0].peer.addr != stale_addr {
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    assert!(
        server.received_routes().is_empty(),
        "stale routes were not flushed"
    );
    Ok(())
}

/// Different routers (different identifiers) never collide.
#[tokio::test]
async fn test_distinct_identifiers_do_not_collide() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut a = TestPeer::connect(server.addr).await?;
    a.handshake(PEER_ASN, HIGH_ID, 90).await?;
    let mut b = TestPeer::connect(server.addr).await?;
    b.handshake(PEER_ASN, LOW_ID, 90).await?;
    wait_for_sessions(&server, 2).await?;
    Ok(())
}

/// A failed connection frees its slot: the same router can come back.
#[tokio::test]
async fn test_identifier_is_released_when_connection_ends() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut first = TestPeer::connect(server.addr).await?;
    first.handshake(PEER_ASN, HIGH_ID, 90).await?;
    wait_for_sessions(&server, 1).await?;
    drop(first);
    wait_for_sessions(&server, 0).await?;

    let mut second = TestPeer::connect(server.addr).await?;
    second.handshake(PEER_ASN, HIGH_ID, 90).await?;
    wait_for_sessions(&server, 1).await?;
    Ok(())
}

/// A peer claiming our own identifier is refused outright.
#[tokio::test]
async fn test_peer_with_our_identifier_is_refused() -> Result<()> {
    let server = TestServer::start(LOCAL_ASN).await?;
    let mut peer = TestPeer::connect(server.addr).await?;
    open_only(&mut peer, Ipv4Addr::new(127, 0, 0, 1)).await?;
    let notif = expect_notification(&mut peer).await?;
    assert!(
        matches!(notif.error, BgpError::OpenMsg(OpenMsgSubcode::BadBgpId)),
        "{notif:?}"
    );
    Ok(())
}