//! Registry of established BGP sessions.
//!
//! A session registers itself when it reaches Established and is removed when
//! its [`SessionGuard`] drops, so a session that errors out is always
//! deregistered. Other parts of the server (the advertiser, the HTTP API) use
//! the registry to enumerate live sessions and queue outbound messages.

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime};

use bgp_packet::errors::{BgpError, CeaseSubcode};
use bgp_packet::message::NotificationMessage;

use bgp_packet::message::Message;
use bgp_primitives::ip_prefix::Nlri;
use bgp_rib::table::PeerIndex;
use tokio::sync::{broadcast, mpsc};

use crate::address::{Endpoint, PeerAddress, socket_text};
use crate::collision::Direction;
use crate::family::FamilySet;
use crate::fsm::Negotiated;
use crate::metrics::{PEER_UP, SESSIONS_ESTABLISHED};
use crate::policy::PeerRole;
use crate::rpki::Subject;
use crate::sync::IgnorePoisonRw;

const EVENT_CHANNEL_CAPACITY: usize = 256;

/// Identifies a session by the peer's TCP address and its AS number.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct PeerId {
    pub addr: SocketAddr,
    pub asn: u32,
}

impl std::fmt::Display for PeerId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}(AS{})", socket_text(self.addr), self.asn)
    }
}

/// Where the dialer is in bringing up a configured peer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DialState {
    /// Not dialling: the peer connected to us, or the dialer has not started.
    Idle,
    Connecting,
    /// The last attempt failed or the session ended; waiting to retry.
    Backoff {
        until: SystemTime,
        /// How the last attempt or session ended.
        last_error: String,
    },
}

/// A peer from the configuration file, whether or not it is currently up.
#[derive(Debug, Clone)]
pub struct ConfiguredPeer {
    pub name: String,
    pub endpoint: Endpoint,
    pub remote_asn: u32,
    /// What the peer is to us, if that was said.
    pub role: Option<PeerRole>,
}

/// Counts kept by a session as it goes, for whoever wants to look.
#[derive(Debug, Default)]
pub struct SessionCounters {
    /// Routes of the peer's that are in the routing table now.
    routes: AtomicU64,
    /// Routes the peer announced that its import policy refused, since the
    /// session came up.
    rejected: AtomicU64,
    /// Routes of the peer's that are held out of the table as invalid, to
    /// be taken if route origin validation comes to say otherwise.
    held: AtomicU64,
}

impl SessionCounters {
    pub fn routes(&self) -> u64 {
        self.routes.load(Ordering::Relaxed)
    }

    pub fn rejected(&self) -> u64 {
        self.rejected.load(Ordering::Relaxed)
    }

    pub fn held(&self) -> u64 {
        self.held.load(Ordering::Relaxed)
    }

    pub(crate) fn set_held(&self, held: usize) {
        self.held.store(held as u64, Ordering::Relaxed);
    }

    /// Notes what one UPDATE did, and returns how many routes the peer has
    /// in the table after it.
    pub(crate) fn note(&self, added: usize, removed: usize, rejected: usize) -> u64 {
        self.rejected.fetch_add(rejected as u64, Ordering::Relaxed);
        // Only the session writes this, so the two steps are not a race.
        let routes = (self.routes() + added as u64).saturating_sub(removed as u64);
        self.routes.store(routes, Ordering::Relaxed);
        routes
    }
}

/// Handle to one established session. Cloning is cheap.
#[derive(Debug, Clone)]
pub struct SessionHandle {
    pub peer: PeerId,
    /// Name from the configuration, if this peer is configured.
    pub name: Option<String>,
    /// Which side opened the TCP connection.
    pub direction: Direction,
    /// Our side of the TCP connection; used as the next hop for routes we
    /// originate to this peer.
    pub local_addr: SocketAddr,
    pub established_at: SystemTime,
    /// Whether the peer is in our own AS.
    pub ibgp: bool,
    /// What the peer is to us, if that is known.
    pub role: Option<PeerRole>,
    /// Whether the peer said so itself in its OPEN (RFC 9234).
    pub role_confirmed: bool,
    /// Whether IPv4 routes go with IPv6 next hops on this session
    /// (RFC 8950).
    pub extended_next_hop: bool,
    /// Whether the peer negotiated RFC 6793 four-octet ASNs. Determines how
    /// AS_PATH is encoded on the wire for this session.
    pub four_octet_asn: bool,
    /// The address families negotiated with the peer. Routes of any other
    /// family are neither sent to it nor accepted from it.
    pub families: FamilySet,
    /// The families for which the peer sends us several routes for a
    /// prefix (ADD-PATH, RFC 7911), and those for which we send it path
    /// identifiers.
    pub add_path_receive: FamilySet,
    pub add_path_send: FamilySet,
    /// The peer's number in the routing table: what its routes there are
    /// marked with.
    pub rib_peer: PeerIndex,
    /// How many routes the peer has, and how many were refused.
    pub counters: Arc<SessionCounters>,
    outbound: mpsc::Sender<Message>,
    /// Which of the RPKI's checks the peer's routes are subject to, if
    /// any, and where the session is told to look at its routes again.
    validation: Option<(Subject, mpsc::UnboundedSender<Vec<Nlri>>)>,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("session with {0} has closed")]
pub struct SessionClosed(pub PeerId);

impl SessionHandle {
    pub(crate) fn new(
        peer: PeerId,
        name: Option<String>,
        direction: Direction,
        local_addr: SocketAddr,
        negotiated: &Negotiated,
        rib_peer: PeerIndex,
        outbound: mpsc::Sender<Message>,
    ) -> Self {
        Self {
            peer,
            name,
            direction,
            local_addr,
            established_at: SystemTime::now(),
            ibgp: negotiated.ibgp,
            role: None,
            role_confirmed: false,
            extended_next_hop: negotiated.extended_next_hop,
            four_octet_asn: negotiated.four_octet_asn,
            families: negotiated.families.clone(),
            add_path_receive: negotiated.add_path_receive.clone(),
            add_path_send: negotiated.add_path_send.clone(),
            rib_peer,
            counters: Arc::default(),
            outbound,
            validation: None,
        }
    }

    /// Marks the session's routes as subject to validation with the
    /// RPKI's data. What is sent with [`SessionHandle::revalidate`]
    /// arrives at the receiver returned.
    pub(crate) fn validate(&mut self, subject: Subject) -> mpsc::UnboundedReceiver<Vec<Nlri>> {
        let (tx, rx) = mpsc::unbounded_channel();
        self.validation = Some((subject, tx));
        rx
    }

    /// Which of the RPKI's checks the session's routes are subject to, if
    /// any.
    pub fn validated(&self) -> Option<Subject> {
        self.validation.as_ref().map(|(subject, _)| *subject)
    }

    /// Tells the session that the validation data has changed. `suspects`
    /// are routes of its peer's in the table that are now invalid.
    pub(crate) fn revalidate(&self, suspects: Vec<Nlri>) {
        if let Some((_, revalidate)) = &self.validation {
            // Fails only if the session has just ended.
            let _ = revalidate.send(suspects);
        }
    }

    /// Whether `other` is a handle to this same session, and not to another
    /// session with the same peer.
    pub fn is_same_session(&self, other: &SessionHandle) -> bool {
        self.outbound.same_channel(&other.outbound)
    }

    /// Queues a message for transmission on this session. Waits for queue
    /// space if the session's writer is behind.
    pub async fn send(&self, msg: Message) -> Result<(), SessionClosed> {
        self.outbound
            .send(msg)
            .await
            .map_err(|_| SessionClosed(self.peer))
    }
}

#[derive(Debug, Clone)]
pub enum SessionEvent {
    Up(SessionHandle),
    Down(PeerId),
    /// The peer sent ROUTE-REFRESH (RFC 2918): it wants everything we
    /// advertise to it sent again, typically because its import policy
    /// changed.
    RefreshRequested(SessionHandle),
}

#[derive(Clone)]
pub struct SessionRegistry {
    inner: Arc<RwLock<HashMap<PeerId, SessionHandle>>>,
    configured: Arc<Vec<ConfiguredPeer>>,
    dial_states: Arc<RwLock<HashMap<String, DialState>>>,
    shutting_down: Arc<AtomicBool>,
    events: broadcast::Sender<SessionEvent>,
}

impl Default for SessionRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl SessionRegistry {
    pub fn new() -> Self {
        Self::with_configured(vec![])
    }

    pub fn with_configured(configured: Vec<ConfiguredPeer>) -> Self {
        let (events, _) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
        let dial_states = configured
            .iter()
            .map(|p| (p.name.clone(), DialState::Idle))
            .collect();
        for peer in &configured {
            PEER_UP.with_label_values(&[&peer.name]).set(0);
        }
        Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
            configured: Arc::new(configured),
            dial_states: Arc::new(RwLock::new(dial_states)),
            shutting_down: Arc::new(AtomicBool::new(false)),
            events,
        }
    }

    /// Peers from the configuration file, each with the dialer's current
    /// state for it.
    pub fn configured_peers(&self) -> Vec<(ConfiguredPeer, DialState)> {
        let states = self.dial_states.reading();
        self.configured
            .iter()
            .map(|p| {
                (
                    p.clone(),
                    states.get(&p.name).cloned().unwrap_or(DialState::Idle),
                )
            })
            .collect()
    }

    pub(crate) fn set_dial_state(&self, name: &str, state: DialState) {
        self.dial_states.writing().insert(name.to_string(), state);
    }

    /// An established session with the peer at this address and with this
    /// ASN, regardless of which side opened the TCP connection.
    pub fn find(&self, address: PeerAddress, asn: u32) -> Option<SessionHandle> {
        self.inner
            .reading()
            .values()
            .find(|h| PeerAddress::from(h.peer.addr) == address && h.peer.asn == asn)
            .cloned()
    }

    /// Adds a session and emits [`SessionEvent::Up`]. The returned guard
    /// removes the session again when dropped.
    pub(crate) fn register(&self, handle: SessionHandle) -> SessionGuard {
        let peer = handle.peer;
        self.inner.writing().insert(peer, handle.clone());
        SESSIONS_ESTABLISHED.inc();
        if let Some(name) = &handle.name {
            PEER_UP.with_label_values(&[name]).set(1);
        }
        let _ = self.events.send(SessionEvent::Up(handle));
        SessionGuard {
            registry: self.clone(),
            peer,
        }
    }

    /// All currently established sessions, in no particular order.
    pub fn list(&self) -> Vec<SessionHandle> {
        self.inner.reading().values().cloned().collect()
    }

    /// Called by a session when its peer asks for a route refresh.
    pub(crate) fn request_refresh(&self, peer: &PeerId) {
        let handle = self.inner.reading().get(peer).cloned();
        if let Some(handle) = handle {
            let _ = self.events.send(SessionEvent::RefreshRequested(handle));
        }
    }

    /// True once [`shutdown`](Self::shutdown) has begun; dialers stop
    /// reconnecting.
    pub fn is_shutting_down(&self) -> bool {
        self.shutting_down.load(Ordering::Relaxed)
    }

    /// Closes every session politely: each peer is sent Cease /
    /// Administrative Shutdown (RFC 4486) so it can drop the session at once
    /// instead of discovering a dead TCP connection. Waits up to `grace` for
    /// the sessions to finish; returns how many were still open.
    pub async fn shutdown(&self, grace: Duration) -> usize {
        self.shutting_down.store(true, Ordering::Relaxed);
        for session in self.list() {
            // Not waited for: a session whose queue is full is one whose
            // peer is not reading, and it would hold up the goodbye to all
            // the others. It ends when the process does.
            let _ = session
                .outbound
                .try_send(Message::Notification(NotificationMessage {
                    error: BgpError::Cease(CeaseSubcode::AdminShutdown),
                    data: vec![],
                }));
        }
        let deadline = tokio::time::Instant::now() + grace;
        while !self.list().is_empty() && tokio::time::Instant::now() < deadline {
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
        self.list().len()
    }

    pub fn subscribe(&self) -> broadcast::Receiver<SessionEvent> {
        self.events.subscribe()
    }

    fn remove(&self, peer: &PeerId) {
        let removed = self.inner.writing().remove(peer);
        if let Some(handle) = &removed {
            SESSIONS_ESTABLISHED.dec();
            if let Some(name) = &handle.name {
                PEER_UP.with_label_values(&[name]).set(0);
            }
            let _ = self.events.send(SessionEvent::Down(*peer));
        }
    }
}

/// Removes the session from its registry on drop.
pub struct SessionGuard {
    registry: SessionRegistry,
    peer: PeerId,
}

impl Drop for SessionGuard {
    fn drop(&mut self) {
        self.registry.remove(&self.peer);
    }
}