use std::convert::Infallible;
use std::time::Duration;

use axum::response::sse::{Event, Sse};
use futures::stream::Stream;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;

use crate::session::SseEvent;

/// Create an SSE stream from an mpsc receiver with disconnect detection.
///
/// - Sends a priming event as the first event (SHOULD per MCP spec 2025-11-25)
/// - Uses CancellationToken for graceful shutdown
/// - Uses `biased` select to prioritize shutdown over messages
/// - KeepAlive sends SSE comments (not events) to avoid confusing MCP clients
pub fn sse_stream(
    rx: mpsc::Receiver<SseEvent>,
    cancel: CancellationToken,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    sse_stream_with_heartbeat_interval(rx, cancel, Duration::from_secs(15), || {})
}

pub fn sse_stream_with_heartbeat<F>(
    rx: mpsc::Receiver<SseEvent>,
    cancel: CancellationToken,
    mut on_keepalive: F,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>>
where
    F: FnMut() + Send + 'static,
{
    sse_stream_with_heartbeat_interval(rx, cancel, Duration::from_secs(15), move || {
        on_keepalive();
    })
}

fn sse_stream_with_heartbeat_interval<F>(
    rx: mpsc::Receiver<SseEvent>,
    cancel: CancellationToken,
    keepalive_interval: Duration,
    mut on_keepalive: F,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>>
where
    F: FnMut() + Send + 'static,
{
    let stream = async_stream::stream! {
        // SSE priming event (SHOULD per spec 2025-11-25):
        // Empty data with event ID so clients know the stream is alive.
        yield Ok(Event::default().id("0").data(""));

        let mut rx = ReceiverStream::new(rx);
        use futures::StreamExt;

        let mut keepalive = tokio::time::interval(keepalive_interval);
        keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        keepalive.tick().await;
        loop {
            tokio::select! {
                biased;
                _ = cancel.cancelled() => {
                    tracing::debug!("SSE stream cancelled via shutdown token");
                    break;
                }
                msg = rx.next() => {
                    match msg {
                        Some(msg) => {
                            yield Ok(Event::default().id(msg.id.to_string()).data(msg.message.as_str()));
                        }
                        None => {
                            tracing::debug!("SSE sender dropped, closing stream");
                            break;
                        }
                    }
                }
                _ = keepalive.tick() => {
                    on_keepalive();
                    yield Ok(Event::default().comment(""));
                }
            }
        }
    };

    Sse::new(stream)
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::response::IntoResponse;

    /// Collect SSE events from the response body until the stream ends or timeout.
    async fn collect_sse_events(body: Body, max_events: usize) -> Vec<String> {
        let mut events = Vec::new();
        let mut stream = body.into_data_stream();
        use futures::StreamExt;

        let timeout = tokio::time::timeout(Duration::from_secs(2), async {
            while let Some(Ok(chunk)) = stream.next().await {
                let text = String::from_utf8_lossy(&chunk).to_string();
                // Split into individual events (separated by double newlines)
                for part in text.split("\n\n") {
                    let trimmed = part.trim();
                    if !trimmed.is_empty() {
                        events.push(trimmed.to_string());
                    }
                }
                if events.len() >= max_events {
                    break;
                }
            }
        });

        let _ = timeout.await;
        events
    }

    #[tokio::test]
    async fn priming_event_is_first() {
        let (tx, rx) = mpsc::channel(8);
        let cancel = CancellationToken::new();

        let sse = sse_stream(rx, cancel.clone());
        let response = sse.into_response();
        let body = response.into_body();

        // Drop sender so stream ends after priming event
        drop(tx);

        let events = collect_sse_events(body, 2).await;
        assert!(!events.is_empty(), "expected at least the priming event");
        // Priming event should have id: 0
        assert!(
            events[0].contains("id: 0"),
            "first event should have id 0, got: {}",
            events[0]
        );
    }

    #[tokio::test]
    async fn messages_forwarded_with_incrementing_ids() {
        let (tx, rx) = mpsc::channel(8);
        let cancel = CancellationToken::new();

        let sse = sse_stream(rx, cancel.clone());
        let response = sse.into_response();
        let body = response.into_body();

        // Send two messages
        tx.send(SseEvent {
            id: 1,
            message: crate::session::SseMessage::from_json_value(
                serde_json::json!({"msg": "hello"}),
            )
            .unwrap(),
        })
        .await
        .unwrap();
        tx.send(SseEvent {
            id: 2,
            message: crate::session::SseMessage::from_json_value(
                serde_json::json!({"msg": "world"}),
            )
            .unwrap(),
        })
        .await
        .unwrap();
        drop(tx);

        let events = collect_sse_events(body, 4).await;
        // Should have: priming (id 0) + 2 messages (id 1, id 2)
        assert!(events.len() >= 3, "expected 3 events, got: {events:?}");
        assert!(events[1].contains("id: 1"), "second event should be id 1");
        assert!(events[2].contains("id: 2"), "third event should be id 2");
        assert!(
            events[1].contains("hello"),
            "event 1 should contain message"
        );
    }

    #[tokio::test]
    async fn cancellation_closes_stream() {
        let (_tx, rx) = mpsc::channel::<SseEvent>(8);
        let cancel = CancellationToken::new();

        let sse = sse_stream(rx, cancel.clone());
        let response = sse.into_response();
        let body = response.into_body();

        // Cancel immediately after priming
        tokio::task::yield_now().await;
        cancel.cancel();

        let events = collect_sse_events(body, 5).await;
        // Should have at most the priming event — stream should close
        assert!(
            events.len() <= 2,
            "stream should close after cancel, got {} events",
            events.len()
        );
    }

    #[tokio::test]
    async fn heartbeat_emits_comment_and_runs_callback() {
        use std::sync::atomic::{AtomicUsize, Ordering};
        let (_tx, rx) = mpsc::channel::<SseEvent>(8);
        let cancel = CancellationToken::new();
        let heartbeats = std::sync::Arc::new(AtomicUsize::new(0));
        let seen = std::sync::Arc::clone(&heartbeats);

        let sse = sse_stream_with_heartbeat_interval(
            rx,
            cancel.clone(),
            Duration::from_millis(10),
            move || {
                seen.fetch_add(1, Ordering::SeqCst);
            },
        );
        let response = sse.into_response();
        let body = response.into_body();

        let events = collect_sse_events(body, 2).await;
        cancel.cancel();

        assert!(heartbeats.load(Ordering::SeqCst) >= 1);
        assert!(
            events.iter().any(|event| event.starts_with(":")),
            "expected at least one keepalive comment event, got: {events:?}"
        );
    }
}
