//! In-process fan-out.
//!
//! One [`tokio::sync::broadcast`] channel per room, created when the first
//! listener arrives and pruned when the last one leaves. This is correct for the
//! deployment both consumers actually have (a single process on a single box) and
//! is knowingly the wrong shape for more than one node. Going multi-node means
//! putting Postgres `LISTEN/NOTIFY` or a message bus behind [`Hub::publish`]; it
//! is a second entry on the same list as the in-process rate limiter, not a
//! surprise.
//!
//! # Why the connection caps are load-bearing
//!
//! Both units run under `MemoryMax=512M`, a hard cgroup cap. Each listener holds
//! a task, a buffer, and a broadcast receiver, and an OOM does not degrade chat,
//! it restarts the whole site. So the caps here are a availability control, not
//! tidiness, and they are enforced before any allocation the connection would
//! own.
//!
//! # Divergence from the SyncKit SSE pattern
//!
//! This is modeled on `server/src/routes/synckit/subscribe.rs` and deviates from
//! it twice, both deliberately:
//!
//! 1. **Pruning is atomic.** SyncKit reads `receiver_count()`, drops the guard,
//!    then calls `remove()`. Between those two steps a new subscriber can take
//!    the entry, and the `remove()` then drops a channel that has a live
//!    listener; the next publish lazily creates a *different* channel and the
//!    listener silently receives nothing. SyncKit tolerates this because its
//!    clients re-pull on their own schedule. A chat client would just sit in a
//!    dead room. [`DashMap::remove_if`] runs the predicate under the shard lock,
//!    which closes the race and still avoids the same-shard deadlock that made
//!    SyncKit split the read and the remove in the first place.
//! 2. **Lag is surfaced, not swallowed.** SyncKit discards `Lagged`. Here it
//!    becomes [`ChatEvent::Gap`] so the client re-fetches from its cursor,
//!    because a room with a silent hole in it is worse than a visible one.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use dashmap::DashMap;
use tokio::sync::broadcast;

use crate::error::{ChatError, LimitScope};
use crate::event::ChatEvent;
use crate::ids::{RoomId, UserId};

/// Caps and buffer sizes. Every consumer sets these explicitly; there is no
/// unbounded configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HubLimits {
    /// Concurrent listeners one user may hold. A few, for multiple tabs. Not a
    /// spam control, a fairness one: without it a single client can consume the
    /// global budget.
    pub max_connections_per_user: usize,
    /// Concurrent listeners across every room in the process. This is the number
    /// that keeps the box alive.
    pub max_connections_total: usize,
    /// Messages buffered per room before a slow listener starts missing events
    /// and gets a [`ChatEvent::Gap`].
    pub room_buffer: usize,
}

impl Default for HubLimits {
    fn default() -> Self {
        Self {
            max_connections_per_user: 4,
            max_connections_total: 2_000,
            room_buffer: 256,
        }
    }
}

struct Inner {
    rooms: DashMap<RoomId, broadcast::Sender<ChatEvent>>,
    per_user: DashMap<UserId, AtomicUsize>,
    total: AtomicUsize,
    limits: HubLimits,
}

/// Per-room fan-out with connection accounting.
///
/// Cheap to clone; both consumers hang one off application state.
#[derive(Clone)]
pub struct Hub {
    inner: Arc<Inner>,
}

impl Hub {
    pub fn new(limits: HubLimits) -> Self {
        Self {
            inner: Arc::new(Inner {
                rooms: DashMap::new(),
                per_user: DashMap::new(),
                total: AtomicUsize::new(0),
                limits,
            }),
        }
    }

    /// Send to everyone listening to `room`, returning how many received it.
    ///
    /// Zero is normal and not an error: an empty room is the common case, and a
    /// message still lands in storage regardless of who is connected.
    pub fn publish(&self, room: RoomId, event: ChatEvent) -> usize {
        self.inner
            .rooms
            .get(&room)
            .and_then(|tx| tx.send(event).ok())
            .unwrap_or(0)
    }

    /// Start listening to `room`.
    ///
    /// # Errors
    ///
    /// [`ChatError::ConnectionLimit`] when the user or the process is at capacity.
    /// The caller has already passed authz by this point, so this is a resource
    /// refusal and should read as one: it is a 503, not a 403.
    pub fn subscribe(&self, room: RoomId, user: UserId) -> Result<Subscription, ChatError> {
        // The permit is acquired first and owns the rollback. There is no window
        // in which a counter is incremented but nothing is responsible for
        // decrementing it, including if anything below panics.
        let permit = ConnectionPermit::acquire(self.inner.clone(), user)?;

        let rx = self
            .inner
            .rooms
            .entry(room)
            .or_insert_with(|| broadcast::channel(self.inner.limits.room_buffer).0)
            .subscribe();

        Ok(Subscription { rx, room, permit })
    }

    /// Rooms currently holding a channel. Diagnostics only.
    pub fn room_count(&self) -> usize {
        self.inner.rooms.len()
    }

    /// Listeners across every room. Worth exposing in a health payload: it is the
    /// number that predicts the memory cap.
    pub fn connection_count(&self) -> usize {
        self.inner.total.load(Ordering::Acquire)
    }

    /// Listeners in one room.
    pub fn room_connection_count(&self, room: RoomId) -> usize {
        self.inner
            .rooms
            .get(&room)
            .map_or(0, |tx| tx.receiver_count())
    }
}

/// Holds one connection's slot in the global and per-user counters, releasing
/// both on drop.
struct ConnectionPermit {
    inner: Arc<Inner>,
    user: UserId,
}

impl ConnectionPermit {
    fn acquire(inner: Arc<Inner>, user: UserId) -> Result<Self, ChatError> {
        // Global first: it is a single atomic, so the cheap check rejects the
        // overload case without touching the per-user map at all.
        let prev = inner.total.fetch_add(1, Ordering::AcqRel);
        if prev >= inner.limits.max_connections_total {
            inner.total.fetch_sub(1, Ordering::AcqRel);
            return Err(ChatError::ConnectionLimit(LimitScope::Global));
        }

        let over_user_limit = {
            let entry = inner
                .per_user
                .entry(user)
                .or_insert_with(|| AtomicUsize::new(0));
            let prev = entry.value().fetch_add(1, Ordering::AcqRel);
            if prev >= inner.limits.max_connections_per_user {
                entry.value().fetch_sub(1, Ordering::AcqRel);
                true
            } else {
                false
            }
        };

        if over_user_limit {
            inner.total.fetch_sub(1, Ordering::AcqRel);
            // The entry may now be at zero and prunable, but leaving it costs one
            // small map slot and the next successful subscribe reuses it. Do not
            // remove here: this path is reached under contention, which is
            // exactly when an extra shard lock is least welcome.
            return Err(ChatError::ConnectionLimit(LimitScope::User));
        }

        Ok(Self { inner, user })
    }
}

impl Drop for ConnectionPermit {
    fn drop(&mut self) {
        self.inner.total.fetch_sub(1, Ordering::AcqRel);

        if let Some(entry) = self.inner.per_user.get(&self.user) {
            entry.value().fetch_sub(1, Ordering::AcqRel);
        }
        // Predicate runs under the shard lock, so a user reconnecting in this
        // instant either keeps their entry or gets a fresh one. Never both.
        self.inner
            .per_user
            .remove_if(&self.user, |_, count| count.load(Ordering::Acquire) == 0);
    }
}

/// A live listener on one room.
///
/// Dropping it releases the connection slot and prunes the room channel if this
/// was the last listener.
pub struct Subscription {
    rx: broadcast::Receiver<ChatEvent>,
    room: RoomId,
    permit: ConnectionPermit,
}

impl Subscription {
    /// Next event, or `None` once the room is gone and nothing more will arrive.
    ///
    /// A listener that fell behind gets [`ChatEvent::Gap`] rather than silence.
    /// The events it missed are not recoverable from the channel by design: the
    /// client re-fetches from its cursor, which is the same path a reconnect
    /// after a deploy takes.
    pub async fn recv(&mut self) -> Option<ChatEvent> {
        match self.rx.recv().await {
            Ok(event) => Some(event),
            Err(broadcast::error::RecvError::Lagged(_)) => Some(ChatEvent::Gap),
            Err(broadcast::error::RecvError::Closed) => None,
        }
    }

    pub fn room(&self) -> RoomId {
        self.room
    }
}

impl Drop for Subscription {
    fn drop(&mut self) {
        // `self.rx` is still alive here (fields drop after this body), so this
        // subscription is still counted. Compare against 1, not 0.
        //
        // The predicate runs under the shard lock, so a subscriber arriving in
        // this instant either keeps this channel or creates a fresh one after the
        // removal. It cannot end up holding a receiver on a sender that has just
        // been dropped from the map, which is the silent-dead-room failure.
        self.permit
            .inner
            .rooms
            .remove_if(&self.room, |_, tx| tx.receiver_count() <= 1);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ids::MessageId;
    use uuid::Uuid;

    fn hub(limits: HubLimits) -> Hub {
        Hub::new(limits)
    }

    fn room() -> RoomId {
        RoomId(Uuid::new_v4())
    }

    fn user() -> UserId {
        UserId(Uuid::new_v4())
    }

    fn delete(id: i64) -> ChatEvent {
        ChatEvent::Delete { id: MessageId(id) }
    }

    #[tokio::test]
    async fn delivers_to_every_listener_in_the_room() {
        let hub = hub(HubLimits::default());
        let r = room();
        let mut a = hub.subscribe(r, user()).unwrap();
        let mut b = hub.subscribe(r, user()).unwrap();

        assert_eq!(hub.publish(r, delete(1)), 2);
        assert_eq!(a.recv().await, Some(delete(1)));
        assert_eq!(b.recv().await, Some(delete(1)));
    }

    #[tokio::test]
    async fn rooms_are_isolated() {
        let hub = hub(HubLimits::default());
        let (r1, r2) = (room(), room());
        let mut listener = hub.subscribe(r1, user()).unwrap();

        assert_eq!(hub.publish(r2, delete(1)), 0);
        assert_eq!(hub.publish(r1, delete(2)), 1);
        assert_eq!(listener.recv().await, Some(delete(2)));
    }

    #[test]
    fn publishing_to_an_empty_room_is_not_an_error() {
        let hub = hub(HubLimits::default());
        assert_eq!(hub.publish(room(), delete(1)), 0);
        // and does not conjure a channel for a room nobody is listening to
        assert_eq!(hub.room_count(), 0);
    }

    #[test]
    fn room_channel_is_pruned_when_the_last_listener_leaves() {
        let hub = hub(HubLimits::default());
        let r = room();

        let a = hub.subscribe(r, user()).unwrap();
        let b = hub.subscribe(r, user()).unwrap();
        assert_eq!(hub.room_count(), 1);

        drop(a);
        assert_eq!(hub.room_count(), 1, "one listener remains");

        drop(b);
        assert_eq!(hub.room_count(), 0, "last listener pruned the room");
    }

    #[tokio::test]
    async fn a_room_recreated_after_pruning_still_delivers() {
        // The failure this guards: pruning a channel that a new subscriber has
        // just taken, leaving them attached to an orphaned sender.
        let hub = hub(HubLimits::default());
        let r = room();

        let first = hub.subscribe(r, user()).unwrap();
        drop(first);
        assert_eq!(hub.room_count(), 0);

        let mut second = hub.subscribe(r, user()).unwrap();
        assert_eq!(hub.publish(r, delete(9)), 1);
        assert_eq!(second.recv().await, Some(delete(9)));
    }

    #[test]
    fn per_user_limit_is_enforced_and_released() {
        let hub = hub(HubLimits {
            max_connections_per_user: 2,
            ..HubLimits::default()
        });
        let r = room();
        let u = user();

        let a = hub.subscribe(r, u).unwrap();
        let b = hub.subscribe(r, u).unwrap();
        assert!(matches!(
            hub.subscribe(r, u),
            Err(ChatError::ConnectionLimit(LimitScope::User))
        ));

        // A different user is unaffected by another's exhausted budget.
        let other = hub.subscribe(r, user()).unwrap();

        drop(a);
        let _c = hub.subscribe(r, u).expect("slot released on drop");
        drop((b, other));
    }

    #[test]
    fn global_limit_is_enforced_and_released() {
        let hub = hub(HubLimits {
            max_connections_total: 2,
            ..HubLimits::default()
        });
        let r = room();

        let a = hub.subscribe(r, user()).unwrap();
        let b = hub.subscribe(r, user()).unwrap();
        assert!(matches!(
            hub.subscribe(r, user()),
            Err(ChatError::ConnectionLimit(LimitScope::Global))
        ));

        drop(a);
        assert_eq!(hub.connection_count(), 1);
        let _c = hub.subscribe(r, user()).expect("slot released on drop");
        drop(b);
    }

    #[test]
    fn a_rejected_subscribe_leaks_no_slot() {
        let hub = hub(HubLimits {
            max_connections_per_user: 1,
            ..HubLimits::default()
        });
        let r = room();
        let u = user();

        let held = hub.subscribe(r, u).unwrap();
        for _ in 0..10 {
            assert!(hub.subscribe(r, u).is_err());
        }
        // The global counter must not have drifted upward on the rejected path.
        assert_eq!(hub.connection_count(), 1);
        drop(held);
        assert_eq!(hub.connection_count(), 0);
    }

    #[tokio::test]
    async fn a_listener_that_falls_behind_gets_a_gap_not_silence() {
        let hub = hub(HubLimits {
            room_buffer: 2,
            ..HubLimits::default()
        });
        let r = room();
        let mut slow = hub.subscribe(r, user()).unwrap();

        for i in 0..10 {
            hub.publish(r, delete(i));
        }

        assert_eq!(slow.recv().await, Some(ChatEvent::Gap));
        // and the stream keeps working afterward
        hub.publish(r, delete(99));
        let next = slow.recv().await;
        assert!(next.is_some(), "stream survives a gap");
    }

    #[test]
    fn counters_return_to_zero_when_everything_drops() {
        let hub = hub(HubLimits::default());
        let r = room();
        let subs: Vec<_> = (0..5).map(|_| hub.subscribe(r, user()).unwrap()).collect();
        assert_eq!(hub.connection_count(), 5);
        assert_eq!(hub.room_connection_count(r), 5);

        drop(subs);
        assert_eq!(hub.connection_count(), 0);
        assert_eq!(hub.room_count(), 0);
        assert_eq!(hub.room_connection_count(r), 0);
    }
}