//! In-memory WebSocket hub: connection registry, room topics, and event fan-out.
//!
//! Implements [`WsGateway`] for publishing [`ChatEventEnvelope`]s to subscribed clients on
//! [`openapi-wss.yaml`](../../../openapi-wss.yaml) `/wss/gateway`.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use async_trait::async_trait;
use athena_chat::events::ChatEventEnvelope;
use athena_chat::ids::{RoomId, UserId};
use axum::Router;
use tokio::sync::{RwLock, mpsc};
use tracing::instrument;
use uuid::Uuid;

use crate::commands::map_chat_event;
use crate::connection::{ConnectionHandle, RoomScopeKey};
use crate::presence;
use crate::protocol::{PresenceSnapshot, ServerWsEvent};
use crate::room_topics::RoomTopic;

/// Gateway surface for mounting routes and pushing realtime events to connected clients.
#[async_trait]
pub trait WsGateway: Send + Sync {
    fn mount_routes(self: Arc<Self>, router: Router<()>) -> Router<()>;
    async fn publish_room_event(
        &self,
        client_name: &str,
        room_id: RoomId,
        event: ChatEventEnvelope,
    );
    async fn publish_user_event(&self, user_id: UserId, event: ServerWsEvent);
    async fn disconnect_user(&self, user_id: UserId, reason: &str);
    async fn current_presence(&self, client_name: &str, room_id: RoomId) -> PresenceSnapshot;
}

/// Shared connection and subscription state for the chat WebSocket gateway.
pub struct WsHub {
    pub connections: Arc<RwLock<HashMap<Uuid, Arc<ConnectionHandle>>>>,
    pub room_subscriptions: Arc<RwLock<HashMap<RoomScopeKey, HashSet<Uuid>>>>,
    pub user_connections: Arc<RwLock<HashMap<UserId, HashSet<Uuid>>>>,
    pub room_topics: Arc<RwLock<HashMap<RoomScopeKey, Arc<RoomTopic>>>>,
    room_channel_capacity: usize,
}

impl WsHub {
    /// Creates an empty hub with the given per-room broadcast channel capacity.
    pub fn new(room_channel_capacity: usize) -> Self {
        Self {
            connections: Arc::new(RwLock::new(HashMap::new())),
            room_subscriptions: Arc::new(RwLock::new(HashMap::new())),
            user_connections: Arc::new(RwLock::new(HashMap::new())),
            room_topics: Arc::new(RwLock::new(HashMap::new())),
            room_channel_capacity,
        }
    }

    /// Registers a new authenticated connection and indexes it by user id.
    pub async fn register_connection(
        &self,
        user_id: UserId,
        organization_id: String,
        client_name: String,
        outbound: mpsc::Sender<ServerWsEvent>,
    ) -> Arc<ConnectionHandle> {
        let handle = Arc::new(ConnectionHandle {
            connection_id: Uuid::new_v4(),
            user_id: user_id.clone(),
            organization_id,
            client_name,
            outbound,
            subscriptions: RwLock::new(HashMap::new()),
        });

        self.connections
            .write()
            .await
            .insert(handle.connection_id, handle.clone());
        self.user_connections
            .write()
            .await
            .entry(user_id)
            .or_default()
            .insert(handle.connection_id);
        handle
    }

    /// Removes a connection and aborts any active room subscription tasks.
    pub async fn unregister_connection(&self, connection_id: Uuid) {
        let handle = self.connections.write().await.remove(&connection_id);
        if let Some(handle) = handle {
            let mut user_connections = self.user_connections.write().await;
            if let Some(connection_ids) = user_connections.get_mut(&handle.user_id) {
                connection_ids.remove(&connection_id);
                if connection_ids.is_empty() {
                    user_connections.remove(&handle.user_id);
                }
            }
            let room_scopes = self
                .room_subscriptions
                .read()
                .await
                .iter()
                .filter_map(|(room_scope, connection_ids)| {
                    connection_ids
                        .contains(&connection_id)
                        .then_some(room_scope.clone())
                })
                .collect::<Vec<_>>();
            for room_scope in room_scopes {
                if let Some(connection_ids) =
                    self.room_subscriptions.write().await.get_mut(&room_scope)
                {
                    connection_ids.remove(&connection_id);
                }
            }
            let mut subscriptions = handle.subscriptions.write().await;
            for (_, task) in subscriptions.drain() {
                task.abort();
            }
        }
    }

    /// Returns or creates the broadcast topic for a tenant-scoped room.
    pub async fn ensure_room_topic(&self, client_name: &str, room_id: RoomId) -> Arc<RoomTopic> {
        let room_scope = RoomScopeKey::new(client_name, room_id);
        if let Some(topic) = self.room_topics.read().await.get(&room_scope) {
            return topic.clone();
        }

        let mut topics = self.room_topics.write().await;
        topics
            .entry(room_scope)
            .or_insert_with(|| Arc::new(RoomTopic::new(self.room_channel_capacity)))
            .clone()
    }

    /// Records that `connection_id` is subscribed to events for the given room scope.
    pub async fn add_room_subscription(
        &self,
        client_name: &str,
        room_id: RoomId,
        connection_id: Uuid,
    ) {
        self.room_subscriptions
            .write()
            .await
            .entry(RoomScopeKey::new(client_name, room_id))
            .or_default()
            .insert(connection_id);
    }

    /// Drops a connection from a room's subscription set without tearing down the topic.
    pub async fn remove_room_subscription(
        &self,
        client_name: &str,
        room_id: RoomId,
        connection_id: Uuid,
    ) {
        let room_scope = RoomScopeKey::new(client_name, room_id);
        if let Some(connections) = self.room_subscriptions.write().await.get_mut(&room_scope) {
            connections.remove(&connection_id);
        }
    }
}

#[async_trait]
impl WsGateway for WsHub {
    fn mount_routes(self: Arc<Self>, router: Router<()>) -> Router<()> {
        router
    }

    #[instrument(name = "wss.publish_room_event", skip(self, event), fields(client_name = %client_name, room_id = %room_id))]
    async fn publish_room_event(
        &self,
        client_name: &str,
        room_id: RoomId,
        event: ChatEventEnvelope,
    ) {
        let topic = self.ensure_room_topic(client_name, room_id).await;
        if let Some(server_event) = map_chat_event(&event) {
            if let Some(seq) = event.room_seq {
                topic.observe_seq(seq);
            }
            let _ = topic.sender.send(server_event);
        }
    }

    async fn publish_user_event(&self, user_id: UserId, event: ServerWsEvent) {
        let connection_ids = self
            .user_connections
            .read()
            .await
            .get(&user_id)
            .cloned()
            .unwrap_or_default();
        let connections = self.connections.read().await;
        for connection_id in connection_ids {
            if let Some(connection) = connections.get(&connection_id) {
                let _ = connection.outbound.send(event.clone()).await;
            }
        }
    }

    async fn disconnect_user(&self, user_id: UserId, reason: &str) {
        self.publish_user_event(
            user_id,
            ServerWsEvent::Error {
                code: "disconnect".to_string(),
                message: reason.to_string(),
            },
        )
        .await;
    }

    async fn current_presence(&self, client_name: &str, room_id: RoomId) -> PresenceSnapshot {
        let connection_ids = self
            .room_subscriptions
            .read()
            .await
            .get(&RoomScopeKey::new(client_name, room_id))
            .cloned()
            .unwrap_or_default();
        let connections = self.connections.read().await;
        let user_ids = connection_ids
            .into_iter()
            .filter_map(|connection_id| {
                connections
                    .get(&connection_id)
                    .map(|handle| handle.user_id.clone())
            })
            .collect::<Vec<_>>();
        presence::snapshot(room_id, user_ids)
    }
}
