//! Live unread counts for the bell. [`NotifyHub`] keeps one `watch` channel per signed-in user who
//! has an open `GET /notifications/stream`. A write that changes a user's inbox publishes the new
//! count, and each open stream of that user sends it. Nothing polls: a stream waits on its channel.
//! The hub is in-process only, so a second server process would not see the change.
//!
//! A channel with no receivers left is removed, so the map holds only users who are connected.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};

use tokio::sync::watch;

use crate::state::AppState;

/// The users with an open stream, each with the channel that carries its unread count.
pub struct NotifyHub {
    users: Mutex<HashMap<String, watch::Sender<u64>>>,
    /// Set once on server shutdown. Every open stream ends when it sees this.
    closing: watch::Sender<bool>,
}

impl NotifyHub {
    pub fn new() -> Self {
        Self {
            users: Mutex::new(HashMap::new()),
            closing: watch::channel(false).0,
        }
    }

    fn users(&self) -> MutexGuard<'_, HashMap<String, watch::Sender<u64>>> {
        self.users.lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// Opens a subscription to the counts published for `user_id`. The first count is not sent
    /// here: the caller reads it from the store after subscribing, so a change in between is
    /// seen by the subscription and sent again rather than lost.
    pub fn subscribe(self: &Arc<Self>, user_id: &str) -> Subscription {
        let rx = self
            .users()
            .entry(user_id.to_owned())
            .or_insert_with(|| watch::channel(0).0)
            .subscribe();
        Subscription {
            rx: Some(rx),
            closing: self.closing.subscribe(),
            hub: Arc::clone(self),
            user_id: user_id.to_owned(),
        }
    }

    /// True when `user_id` has an open stream. A write skips the count query for anyone else.
    pub fn is_watched(&self, user_id: &str) -> bool {
        self.users()
            .get(user_id)
            .is_some_and(|tx| tx.receiver_count() > 0)
    }

    /// Sends `count` to every open stream of `user_id`. Drops the channel when none is left.
    pub fn publish(&self, user_id: &str, count: u64) {
        let mut users = self.users();
        let Some(tx) = users.get(user_id) else {
            return;
        };
        if tx.receiver_count() == 0 {
            users.remove(user_id);
            return;
        }
        tx.send_replace(count);
    }

    /// Ends every open stream. Called once when the server shuts down.
    pub fn close(&self) {
        self.closing.send_replace(true);
    }

    /// Removes the user's channel once its last subscription is gone.
    fn release(&self, user_id: &str) {
        let mut users = self.users();
        if users
            .get(user_id)
            .is_some_and(|tx| tx.receiver_count() == 0)
        {
            users.remove(user_id);
        }
    }
}

/// One open stream's view of its user's count. Dropping it releases the user's channel when no
/// other stream of that user is open.
pub struct Subscription {
    /// Taken in `drop` before the hub is told, so the hub counts this receiver as gone.
    rx: Option<watch::Receiver<u64>>,
    closing: watch::Receiver<bool>,
    hub: Arc<NotifyHub>,
    user_id: String,
}

impl Subscription {
    /// Waits for the next published count. `None` once the hub is closing (server shutdown).
    pub async fn next_count(&mut self) -> Option<u64> {
        loop {
            if *self.closing.borrow() {
                return None;
            }
            let rx = self.rx.as_mut()?;
            tokio::select! {
                changed = rx.changed() => {
                    changed.ok()?;
                    return Some(*rx.borrow_and_update());
                }
                changed = self.closing.changed() => {
                    changed.ok()?;
                }
            }
        }
    }
}

impl Drop for Subscription {
    fn drop(&mut self) {
        self.rx = None;
        self.hub.release(&self.user_id);
    }
}

/// Publishes `user_id`'s current unread count to that user's open streams. Call it after a write
/// that changes the inbox. Skips the count query when nobody is listening. Never fails: problems
/// are logged.
pub async fn publish_unread(state: &AppState, user_id: &str) {
    if !state.notify_hub.is_watched(user_id) {
        return;
    }
    match state.store.unread_count(user_id).await {
        Ok(count) => state.notify_hub.publish(user_id, count),
        Err(e) => tracing::error!(error = %e, user_id, "notify: unread count failed"),
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;

    #[tokio::test]
    async fn publish_reaches_the_subscriber_of_that_user_only() {
        let hub = Arc::new(NotifyHub::new());
        let mut alice = hub.subscribe("alice");
        let mut bob = hub.subscribe("bob");

        hub.publish("alice", 3);
        assert_eq!(alice.next_count().await, Some(3));

        // Bob heard nothing: his wait times out.
        let quiet = tokio::time::timeout(Duration::from_millis(50), bob.next_count()).await;
        assert!(quiet.is_err());
    }

    #[tokio::test]
    async fn publish_without_subscribers_is_a_no_op_and_nothing_is_kept() {
        let hub = Arc::new(NotifyHub::new());
        assert!(!hub.is_watched("alice"));
        hub.publish("alice", 1);
        assert!(!hub.is_watched("alice"));
        assert!(hub.users().is_empty());
    }

    #[tokio::test]
    async fn the_channel_is_released_when_the_last_subscription_drops() {
        let hub = Arc::new(NotifyHub::new());
        let first = hub.subscribe("alice");
        let second = hub.subscribe("alice");
        assert!(hub.is_watched("alice"));

        drop(first);
        assert!(hub.is_watched("alice"));

        drop(second);
        assert!(!hub.is_watched("alice"));
        assert!(hub.users().is_empty());
    }

    #[tokio::test]
    async fn every_open_stream_of_a_user_gets_the_count() {
        let hub = Arc::new(NotifyHub::new());
        let mut tab_one = hub.subscribe("alice");
        let mut tab_two = hub.subscribe("alice");

        hub.publish("alice", 5);
        assert_eq!(tab_one.next_count().await, Some(5));
        assert_eq!(tab_two.next_count().await, Some(5));
    }

    #[tokio::test]
    async fn close_ends_every_subscription() {
        let hub = Arc::new(NotifyHub::new());
        let mut sub = hub.subscribe("alice");

        hub.close();
        assert_eq!(sub.next_count().await, None);
    }
}