//! A minimal stand-in for a lab router: connects to a ferrix_bgp instance,
//! completes the handshake, optionally announces some prefixes, and prints
//! every UPDATE it receives. Useful for demos when no real router is around.
//!
//!     cargo run -p bgp-integration-tests --example fake_peer -- \
//!         --server 127.0.0.1:1790 --asn 65001 --announce 203.0.113.0/24
//!
//! With `--echo`, every prefix received is re-announced back to the server
//! with this peer's ASN prepended, which lets a single ferrix instance show
//! its own routes "coming back" in the received table.
//!
//! With `--listen ADDR` it instead waits for the server to dial it (one
//! connection at a time), which exercises the server's active side.

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

use bgp_integration_tests::test_peer::TestPeer;
use bgp_packet::message::{
    AsPathAttribute, CommunitiesPathAttribute, Message, NextHopPathAttribute, OriginPathAttribute,
    PathAttribute, UpdateMessage,
};
use bgp_primitives::ip_prefix::IpPrefix;
use clap::Parser;
use eyre::Result;
use tokio::net::TcpListener;

#[derive(Parser, Debug)]
struct Args {
    /// Address of the ferrix_bgp listener to connect to.
    #[arg(long, default_value = "127.0.0.1:1790", conflicts_with = "listen")]
    server: String,
    /// Instead of connecting, listen here and accept connections from the
    /// server, one at a time, forever.
    #[arg(long)]
    listen: Option<String>,
    /// ASN this peer speaks as.
    #[arg(long, default_value_t = 65001)]
    asn: u32,
    /// Router ID to send in OPEN.
    #[arg(long, default_value = "10.0.0.2")]
    router_id: Ipv4Addr,
    /// Hold time to negotiate.
    #[arg(long, default_value_t = 30)]
    hold_time: u16,
    /// Prefixes to announce after the session is established.
    #[arg(long)]
    announce: Vec<String>,
    /// Re-announce every received prefix back to the server, with our ASN
    /// prepended and NO_EXPORT-tagged routes dropped.
    #[arg(long)]
    echo: bool,
}

const NO_EXPORT: (u16, u16) = (0xFFFF, 0xFF01);

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();
    match &args.listen {
        Some(listen) => {
            let listener = TcpListener::bind(listen).await?;
            println!("listening on {}", listener.local_addr()?);
            loop {
                let (stream, from) = listener.accept().await?;
                println!("accepted connection from {from}");
                if let Err(e) = run(&args, TestPeer::from_stream(stream), from).await {
                    println!("session ended: {e}");
                }
            }
        }
        None => {
            let addr = args.server.parse()?;
            run(&args, TestPeer::connect(addr).await?, addr).await
        }
    }
}

async fn run(args: &Args, mut peer: TestPeer, addr: std::net::SocketAddr) -> Result<()> {
    let open = peer.handshake(args.asn, args.router_id, args.hold_time).await?;
    println!("established with {} (AS{} router-id {})", addr, open.asn, open.identifier);

    for p in &args.announce {
        let prefix = IpPrefix::try_from(p.as_str())?;
        let update = announce(prefix.clone(), vec![args.asn], vec![(args.asn as u16, 1)]);
        peer.send(&Message::Update(update)).await?;
        println!("announced {prefix}");
    }

    let keepalive = Duration::from_secs((args.hold_time / 3).max(1) as u64);
    let mut ticker = tokio::time::interval(keepalive);
    loop {
        tokio::select! {
            _ = ticker.tick() => peer.send(&Message::KeepAlive).await?,
            msg = peer.recv() => match msg? {
                Message::KeepAlive => {}
                Message::Update(u) => {
                    for w in &u.withdrawn_routes {
                        println!("WITHDRAW {w}");
                    }
                    for a in &u.announced_routes {
                        println!("UPDATE {a} attrs={:?}", u.path_attributes);
                    }
                    if args.echo {
                        echo_back(&mut peer, &u, args.asn).await?;
                    }
                }
                Message::Notification(n) => {
                    println!("NOTIFICATION {}; session over", n.error);
                    return Ok(());
                }
                other => println!("unexpected {other:?}"),
            },
        }
    }
}

fn announce(prefix: IpPrefix, as_path: Vec<u32>, communities: Vec<(u16, u16)>) -> UpdateMessage {
    let mut path_attributes = vec![
        PathAttribute::Origin(OriginPathAttribute::IGP),
        AsPathAttribute::from_asns(as_path),
        PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(10, 0, 0, 2))),
    ];
    if !communities.is_empty() {
        path_attributes.push(PathAttribute::Communities(CommunitiesPathAttribute(communities)));
    }
    UpdateMessage { withdrawn_routes: vec![], path_attributes, announced_routes: vec![prefix] }
}

async fn echo_back(peer: &mut TestPeer, u: &UpdateMessage, asn: u32) -> Result<()> {
    let mut as_path = vec![asn];
    let mut communities = vec![];
    for attr in &u.path_attributes {
        match attr {
            PathAttribute::AsPath(p) => as_path.extend(p.segments.iter().flat_map(|s| s.path.iter().copied())),
            PathAttribute::Communities(c) => communities = c.0.clone(),
            _ => {}
        }
    }
    if communities.contains(&NO_EXPORT) {
        println!("  not echoing: NO_EXPORT");
        return Ok(());
    }
    for prefix in &u.announced_routes {
        peer.send(&Message::Update(announce(prefix.clone(), as_path.clone(), communities.clone()))).await?;
        println!("  echoed {prefix} back with path {as_path:?}");
    }
    if !u.withdrawn_routes.is_empty() {
        peer.send(&Message::Update(UpdateMessage {
            withdrawn_routes: u.withdrawn_routes.clone(),
            path_attributes: vec![],
            announced_routes: vec![],
        }))
        .await?;
        println!("  echoed withdraw of {:?}", u.withdrawn_routes.iter().map(|p| p.to_string()).collect::<Vec<_>>());
    }
    Ok(())
}