//! SSE push notification endpoint for SyncKit subscribers.
//!
//! Clients connect to `GET /api/sync/subscribe?app_id={uuid}` and receive
//! `changed` events whenever another device pushes changes. The payload is
//! `{"seq":N}`, the new max seq, so a client already at that cursor can skip
//! the pull; everything else is pulled. No content crosses the stream, which
//! preserves E2E encryption (server never sends plaintext).

use std::convert::Infallible;
use std::sync::atomic::Ordering;
use std::time::Duration;

use axum::{
    extract::{Query, State},
    response::sse::{Event, KeepAlive, Sse},
};
use serde::Deserialize;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::BroadcastStream;

use crate::{
    constants,
    db::{SyncAppId, UserId},
    error::{AppError, Result},
    synckit_auth::SyncUser,
};

/// Drop guard that decrements the per-user SSE connection counter.
struct SseConnectionGuard {
    sse_connections: std::sync::Arc<dashmap::DashMap<UserId, std::sync::atomic::AtomicUsize>>,
    sync_notify:
        std::sync::Arc<dashmap::DashMap<(SyncAppId, UserId), tokio::sync::broadcast::Sender<i64>>>,
    user_id: UserId,
    app_id: SyncAppId,
}

impl Drop for SseConnectionGuard {
    fn drop(&mut self) {
        // Decrement the connection counter, then drop the entry if it hit zero.
        // `remove_if` runs the predicate under the shard lock, so a connection
        // arriving concurrently cannot be counted into an entry we then remove.
        // It also avoids the same-shard deadlock a get() + remove() pair would
        // hit, which is what the earlier read/drop/remove split was working
        // around.
        if let Some(counter) = self.sse_connections.get(&self.user_id) {
            counter.value().fetch_sub(1, Ordering::AcqRel);
        }
        self.sse_connections.remove_if(&self.user_id, |_, counter| {
            counter.load(Ordering::Acquire) == 0
        });

        // Prune the sync_notify channel once the last receiver drops. This must
        // be atomic with respect to `Sync::subscribe_channel`: reading
        // receiver_count(), dropping the guard, then removing lets a subscriber
        // land in between and hold a receiver on a sender we then discard. The
        // next push lazily creates a *different* channel, and that client
        // silently receives no further SSE events until it reconnects.
        let key = (self.app_id, self.user_id);
        self.sync_notify
            .remove_if(&key, |_, sender| sender.receiver_count() == 0);
    }
}

/// Stream wrapper that holds an SSE connection guard. When the stream is
/// dropped (client disconnects), the guard decrements the connection counter.
struct GuardedStream<S> {
    inner: S,
    _guard: SseConnectionGuard,
}

impl<S: tokio_stream::Stream + Unpin> tokio_stream::Stream for GuardedStream<S> {
    type Item = S::Item;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        std::pin::Pin::new(&mut self.inner).poll_next(cx)
    }
}

#[derive(Deserialize)]
pub(super) struct SubscribeQuery {
    pub app_id: SyncAppId,
}

/// SSE endpoint for real-time sync push notifications.
///
/// `GET /api/sync/subscribe?app_id={uuid}`: JWT auth required.
///
/// Returns an SSE stream that emits `event: changed` with a `{"seq":N}` payload
/// whenever a push is made to the same app+user. A client behind that seq pulls.
/// A keepalive comment is sent every 30 seconds to prevent proxy timeouts.
#[tracing::instrument(skip_all, name = "synckit::subscribe")]
pub(super) async fn sync_subscribe(
    State(sync): State<crate::Sync>,
    sync_user: SyncUser,
    Query(query): Query<SubscribeQuery>,
) -> Result<Sse<impl tokio_stream::Stream<Item = std::result::Result<Event, Infallible>>>> {
    // Validate that the requested app_id matches the JWT's app_id
    if query.app_id != sync_user.app_id {
        return Err(AppError::BadRequest(
            "app_id does not match authenticated session".to_string(),
        ));
    }

    // Enforce per-user SSE connection limit
    let counter = sync
        .sse_connections
        .entry(sync_user.user_id)
        .or_insert_with(|| std::sync::atomic::AtomicUsize::new(0));
    let current = counter.value().fetch_add(1, Ordering::AcqRel);
    if current >= constants::SYNCKIT_MAX_SSE_CONNECTIONS_PER_USER {
        counter.value().fetch_sub(1, Ordering::AcqRel);
        return Err(AppError::BadRequest(
            "Too many concurrent SSE connections".to_string(),
        ));
    }

    let guard = SseConnectionGuard {
        sse_connections: sync.sse_connections.clone(),
        sync_notify: sync.sync_notify.clone(),
        user_id: sync_user.user_id,
        app_id: sync_user.app_id,
    };

    // Get or create the broadcast channel for this app+user
    let rx = sync.subscribe_channel(sync_user.app_id, sync_user.user_id);

    let stream = BroadcastStream::new(rx).filter_map(|result| match result {
        // Carry the new max seq so a client already at this cursor can skip the
        // pull. Clients that don't read `seq` just pull as before.
        Ok(seq) => Some(Ok(Event::default()
            .event("changed")
            .data(format!("{{\"seq\":{seq}}}")))),
        Err(_) => None, // Lagged, skip missed events, client will pull anyway
    });

    // Wrap stream with the connection guard, when the client disconnects and
    // the stream is dropped, the guard's Drop impl decrements the counter.
    let stream = GuardedStream {
        inner: stream,
        _guard: guard,
    };

    Ok(Sse::new(stream).keep_alive(
        KeepAlive::new()
            .interval(Duration::from_secs(30))
            .text("keepalive"),
    ))
}