//! Per-connection session state for the Streamable HTTP transport.
//!
//! A session is created on `initialize` (the server returns its id in the
//! `Mcp-Session-Id` header) and referenced by that id on every later request.
//! It holds the three things the push layer needs that a stateless request
//! cannot:
//!
//! - the set of resource uris this client has **subscribed** to,
//! - the **out** channel to this client's SSE stream (present while a GET
//!   stream is open), down which the server pushes notifications,
//! - the **in-flight** `tools/call` requests, each with a trigger to cancel it.
//!
//! stdio has no sessions (one implicit peer, no `Mcp-Session-Id`), so its
//! dispatch passes `None` and the session-scoped methods are unavailable there.

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

use serde_json::Value;
use tokio::sync::{mpsc, oneshot};

use crate::error::{Error, Result};
use crate::tool::ToolCallResult;

/// JSON-RPC notification value pushed to a session's SSE stream.
pub(crate) type OutSender = mpsc::UnboundedSender<Value>;

struct Inflight {
    /// Original JSON-RPC id, preserved with its type (number stays a number)
    /// so a `notifications/cancelled` echoes the exact id the client sent.
    id: Value,
    cancel: oneshot::Sender<()>,
}

/// One client connection's mutable state. Shared (`Arc`) between the POST
/// handler that mutates it and the GET SSE handler that drains its out channel.
pub(crate) struct Session {
    pub id: String,
    subscriptions: Mutex<HashSet<String>>,
    out: Mutex<Option<OutSender>>,
    inflight: Mutex<HashMap<String, Inflight>>,
}

impl Session {
    fn new(id: String) -> Self {
        Self {
            id,
            subscriptions: Mutex::new(HashSet::new()),
            out: Mutex::new(None),
            inflight: Mutex::new(HashMap::new()),
        }
    }

    pub(crate) fn subscribe(&self, uri: impl Into<String>) {
        self.subscriptions.lock().unwrap().insert(uri.into());
    }

    /// Remove a subscription. Returns whether it was present.
    pub(crate) fn unsubscribe(&self, uri: &str) -> bool {
        self.subscriptions.lock().unwrap().remove(uri)
    }

    pub(crate) fn is_subscribed(&self, uri: &str) -> bool {
        self.subscriptions.lock().unwrap().contains(uri)
    }

    /// Attach the SSE out channel (called when the client opens GET `/mcp`).
    /// Replaces any prior channel — a reconnect supersedes the stale stream.
    pub(crate) fn attach_out(&self, tx: OutSender) {
        *self.out.lock().unwrap() = Some(tx);
    }

    /// Push a notification to the SSE stream if one is connected.
    pub(crate) fn push(&self, notification: Value) {
        if let Some(tx) = self.out.lock().unwrap().as_ref() {
            let _ = tx.send(notification);
        }
    }

    /// Run a `tools/call` future so it can be cancelled while in flight. The
    /// call is registered under `id` until it completes; a concurrent
    /// [`cancel_request`](Self::cancel_request) or [`cancel_all`](Self::cancel_all)
    /// aborts it with [`Error::Cancelled`].
    pub(crate) async fn run_cancellable<F>(&self, id: &Value, fut: F) -> Result<ToolCallResult>
    where
        F: Future<Output = Result<ToolCallResult>>,
    {
        let key = id.to_string();
        let (tx, rx) = oneshot::channel();
        self.inflight.lock().unwrap().insert(
            key.clone(),
            Inflight {
                id: id.clone(),
                cancel: tx,
            },
        );

        let outcome = tokio::select! {
            biased;
            _ = rx => Err(Error::Cancelled(key.clone())),
            result = fut => result,
        };

        self.inflight.lock().unwrap().remove(&key);
        outcome
    }

    /// Cancel one in-flight request by its (stringified) id. Triggers the
    /// select in [`run_cancellable`](Self::run_cancellable). No-op if the id is
    /// not in flight (already finished or never existed).
    pub(crate) fn cancel_request(&self, id_key: &str) {
        if let Some(entry) = self.inflight.lock().unwrap().remove(id_key) {
            let _ = entry.cancel.send(());
        }
    }

    /// Cancel every in-flight request for this session and tell the client,
    /// pushing a `notifications/cancelled` for each. This is the app-initiated
    /// path — a human taking a buffer back. Returns the count cancelled.
    pub(crate) fn cancel_all(&self, reason: &str) -> usize {
        let drained: Vec<Inflight> = {
            let mut map = self.inflight.lock().unwrap();
            map.drain().map(|(_, v)| v).collect()
        };
        let n = drained.len();
        for entry in drained {
            let _ = entry.cancel.send(());
            self.push(cancelled_notification(&entry.id, reason));
        }
        n
    }
}

/// The set of live sessions, shared by the server's handlers and its
/// [`BoundServer`](super::BoundServer) cancel handle.
#[derive(Clone, Default)]
pub(crate) struct Sessions {
    map: Arc<Mutex<HashMap<String, Arc<Session>>>>,
}

impl Sessions {
    /// Create a fresh session with a random id and return it.
    pub(super) fn create(&self) -> Arc<Session> {
        let id = uuid::Uuid::new_v4().to_string();
        let session = Arc::new(Session::new(id.clone()));
        self.map.lock().unwrap().insert(id, session.clone());
        session
    }

    pub(super) fn get(&self, id: &str) -> Option<Arc<Session>> {
        self.map.lock().unwrap().get(id).cloned()
    }

    /// Remove a session (client sent DELETE `/mcp`). Returns whether present.
    pub(super) fn remove(&self, id: &str) -> bool {
        self.map.lock().unwrap().remove(id).is_some()
    }

    pub(super) fn ids(&self) -> Vec<String> {
        self.map.lock().unwrap().keys().cloned().collect()
    }

    /// Cancel in-flight ops across every session. Returns the total cancelled.
    pub(super) fn cancel_all(&self, reason: &str) -> usize {
        let sessions: Vec<Arc<Session>> = self.map.lock().unwrap().values().cloned().collect();
        sessions.iter().map(|s| s.cancel_all(reason)).sum()
    }
}

/// Build a server→client `notifications/cancelled` for `id`.
pub(crate) fn cancelled_notification(id: &Value, reason: &str) -> Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "method": "notifications/cancelled",
        "params": { "requestId": id, "reason": reason },
    })
}

/// Build a server→client `notifications/resources/updated` for `uri`.
pub(crate) fn resource_updated_notification(uri: &str) -> Value {
    serde_json::json!({
        "jsonrpc": "2.0",
        "method": "notifications/resources/updated",
        "params": { "uri": uri },
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn run_cancellable_returns_result_when_not_cancelled() {
        let s = Session::new("s1".into());
        let out = s
            .run_cancellable(&serde_json::json!(1), async {
                Ok(ToolCallResult::text("done"))
            })
            .await
            .unwrap();
        assert_eq!(out.content.len(), 1);
    }

    #[tokio::test]
    async fn cancel_request_aborts_an_in_flight_call() {
        let s = Arc::new(Session::new("s1".into()));
        let id = serde_json::json!(7);
        let s2 = s.clone();
        // A call that would never finish on its own.
        let never = async { std::future::pending::<Result<ToolCallResult>>().await };
        let handle = tokio::spawn(async move { s2.run_cancellable(&id, never).await });
        // Let the call register itself, then cancel it.
        tokio::task::yield_now().await;
        s.cancel_request("7");
        let outcome = handle.await.unwrap();
        assert!(matches!(outcome, Err(Error::Cancelled(_))));
    }

    #[tokio::test]
    async fn cancel_all_notifies_the_client_with_original_id_type() {
        let s = Arc::new(Session::new("s1".into()));
        let (tx, mut rx) = mpsc::unbounded_channel();
        s.attach_out(tx);

        let id = serde_json::json!(42); // numeric id
        let s2 = s.clone();
        let never = async { std::future::pending::<Result<ToolCallResult>>().await };
        let handle = tokio::spawn(async move { s2.run_cancellable(&id, never).await });
        tokio::task::yield_now().await;

        assert_eq!(s.cancel_all("human took the buffer back"), 1);
        assert!(matches!(handle.await.unwrap(), Err(Error::Cancelled(_))));

        let note = rx.recv().await.unwrap();
        assert_eq!(note["method"], "notifications/cancelled");
        assert_eq!(note["params"]["requestId"], 42); // stayed a JSON number
        assert_eq!(note["params"]["reason"], "human took the buffer back");
    }

    #[test]
    fn subscribe_unsubscribe_tracks_membership() {
        let s = Session::new("s1".into());
        assert!(!s.is_subscribed("buffer://a"));
        s.subscribe("buffer://a");
        assert!(s.is_subscribed("buffer://a"));
        assert!(s.unsubscribe("buffer://a"));
        assert!(!s.unsubscribe("buffer://a"));
    }

    #[test]
    fn sessions_create_get_remove() {
        let sessions = Sessions::default();
        let s = sessions.create();
        assert!(sessions.get(&s.id).is_some());
        assert_eq!(sessions.ids().len(), 1);
        assert!(sessions.remove(&s.id));
        assert!(sessions.get(&s.id).is_none());
    }
}