ahp 0.1.0

//! Async JSON-RPC client.
//!
//! The [`Client`] drives a [`Transport`] through a background task and
//! exposes an ergonomic request/notification API on top of it. The
//! transport is pluggable — WebSocket, stdio, in-memory pairs for tests —
//! anything that can deliver framed JSON-RPC messages works.
//!
//! ## Lifecycle
//!
//! ```no_run
//! # async fn run(transport: impl ahp::Transport) -> Result<(), ahp::ClientError> {
//! use ahp::{Client, ClientConfig};
//!
//! let client = Client::connect(transport, ClientConfig::default()).await?;
//! let init = client.initialize("my-client".into(), 1, vec![]).await?;
//! // ... use client ...
//! client.shutdown().await;
//! # Ok(()) }
//! ```
//!
//! The background task terminates when [`Client::shutdown`] is called, the
//! transport closes, or the last [`Client`] handle is dropped.

use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc,
};
use std::time::Duration;

use ahp_types::actions::{ActionEnvelope, StateAction};
use ahp_types::commands::{
    DispatchActionParams, InitializeParams, InitializeResult, ReconnectParams, ReconnectResult,
    SubscribeParams, SubscribeResult, UnsubscribeParams,
};
use ahp_types::messages::{
    ActionNotificationParams, JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest,
    JsonRpcVersion, NotificationMethodParams,
};
use ahp_types::notifications::ProtocolNotification;
use serde::{de::DeserializeOwned, Serialize};
use serde_json::Value;
use tokio::sync::{broadcast, mpsc, oneshot, Mutex};
use tokio::task::JoinHandle;

use crate::error::ClientError;
use crate::transport::{Transport, TransportMessage};

/// Default size of a per-subscription broadcast channel. Consumers that
/// lag behind this many messages will drop frames and must resubscribe.
const DEFAULT_SUBSCRIPTION_BUFFER: usize = 256;

/// Configuration for a [`Client`].
#[derive(Debug, Clone)]
pub struct ClientConfig {
    /// How long to wait for a request to resolve before failing with
    /// [`ClientError::Cancelled`]. `None` disables the default timeout.
    pub default_request_timeout: Option<Duration>,
    /// Size of each subscription's broadcast ring buffer.
    pub subscription_buffer: usize,
}

impl Default for ClientConfig {
    fn default() -> Self {
        Self {
            default_request_timeout: Some(Duration::from_secs(30)),
            subscription_buffer: DEFAULT_SUBSCRIPTION_BUFFER,
        }
    }
}

/// Inbound event fanned out to a [`SessionSubscription`].
///
/// `Action` envelopes carry the write-ahead mutation stream; `Notification`
/// frames carry protocol-level signals the server broadcasts (e.g. session
/// added/removed, connectivity changes).
#[derive(Debug, Clone)]
pub enum SubscriptionEvent {
    /// A write-ahead action envelope for this subscription's resource.
    Action(ActionEnvelope),
    /// A broadcast notification (shared across all subscriptions).
    Notification(ProtocolNotification),
}

/// Handle to a single resource subscription. Drop to stop receiving
/// events. The underlying server subscription is released when the last
/// handle for that URI is dropped and [`Client::unsubscribe`] is called.
pub struct SessionSubscription {
    rx: broadcast::Receiver<SubscriptionEvent>,
    uri: String,
}

impl SessionSubscription {
    /// Resource URI this subscription is bound to.
    pub fn uri(&self) -> &str {
        &self.uri
    }

    /// Await the next event on this subscription. Returns `None` when the
    /// client has shut down.
    pub async fn recv(&mut self) -> Option<SubscriptionEvent> {
        loop {
            match self.rx.recv().await {
                Ok(ev) => return Some(ev),
                Err(broadcast::error::RecvError::Closed) => return None,
                // Slow consumer: skip the gap and keep going. Callers
                // that need strict ordering should use a tighter buffer
                // or their own backpressure.
                Err(broadcast::error::RecvError::Lagged(_)) => continue,
            }
        }
    }
}

/// Handle returned from [`Client::dispatch`].
///
/// Dispatch is fire-and-forget by design; the handle simply records the
/// `clientSeq` that was assigned so callers can correlate their local
/// reducer state with server echoes.
#[derive(Debug, Clone, Copy)]
pub struct DispatchHandle {
    /// Client-local sequence number assigned to this dispatch.
    pub client_seq: i64,
}

// ─── Internal plumbing ───────────────────────────────────────────────────────

type PendingMap = HashMap<u64, oneshot::Sender<Result<Value, JsonRpcError>>>;

struct Shared {
    pending: Mutex<PendingMap>,
    subscriptions: Mutex<HashMap<String, broadcast::Sender<SubscriptionEvent>>>,
    outbound: mpsc::Sender<Outbound>,
    next_id: AtomicU64,
    next_client_seq: AtomicU64,
    config: ClientConfig,
}

enum Outbound {
    Message(JsonRpcMessage),
    Shutdown,
}

/// Async JSON-RPC client driving a pluggable [`Transport`].
///
/// Cheaply cloneable — all clones share the same background task, pending
/// request map, and subscription registry.
#[derive(Clone)]
pub struct Client {
    shared: Arc<Shared>,
    _reader: Arc<DriveHandle>,
}

struct DriveHandle {
    handle: Mutex<Option<JoinHandle<()>>>,
}

impl Drop for DriveHandle {
    fn drop(&mut self) {
        if let Ok(mut guard) = self.handle.try_lock() {
            if let Some(h) = guard.take() {
                h.abort();
            }
        }
    }
}

impl Client {
    /// Connect a client to the given transport. Spawns a background task
    /// that drives the transport until shutdown.
    pub async fn connect<T: Transport>(
        transport: T,
        config: ClientConfig,
    ) -> Result<Self, ClientError> {
        let (outbound_tx, outbound_rx) = mpsc::channel::<Outbound>(64);
        let shared = Arc::new(Shared {
            pending: Mutex::new(HashMap::new()),
            subscriptions: Mutex::new(HashMap::new()),
            outbound: outbound_tx,
            next_id: AtomicU64::new(1),
            next_client_seq: AtomicU64::new(1),
            config,
        });

        let handle = tokio::spawn(drive_transport(transport, shared.clone(), outbound_rx));
        Ok(Self {
            shared,
            _reader: Arc::new(DriveHandle {
                handle: Mutex::new(Some(handle)),
            }),
        })
    }

    /// Gracefully shut down the client, aborting any in-flight requests
    /// with [`ClientError::Shutdown`].
    pub async fn shutdown(&self) {
        let _ = self.shared.outbound.send(Outbound::Shutdown).await;
        // Fail any pending in-flight requests.
        let mut pending = self.shared.pending.lock().await;
        for (_, tx) in pending.drain() {
            let _ = tx.send(Err(JsonRpcError {
                code: -32000,
                message: "client shut down".into(),
                data: None,
            }));
        }
    }

    /// Send a JSON-RPC request and await its result.
    pub async fn request<P, R>(&self, method: &str, params: P) -> Result<R, ClientError>
    where
        P: Serialize,
        R: DeserializeOwned,
    {
        let id = self.shared.next_id.fetch_add(1, Ordering::Relaxed);
        let params_val = serde_json::to_value(&params)?;
        let params_any = if params_val.is_null() {
            None
        } else {
            Some(ahp_types::common::AnyValue::from(params_val))
        };
        let req = JsonRpcMessage::Request(JsonRpcRequest {
            jsonrpc: JsonRpcVersion::V2,
            id,
            method: method.into(),
            params: params_any,
        });

        let (tx, rx) = oneshot::channel();
        {
            let mut pending = self.shared.pending.lock().await;
            pending.insert(id, tx);
        }

        if self
            .shared
            .outbound
            .send(Outbound::Message(req))
            .await
            .is_err()
        {
            self.shared.pending.lock().await.remove(&id);
            return Err(ClientError::Shutdown);
        }

        let result = match self.shared.config.default_request_timeout {
            Some(dur) => match tokio::time::timeout(dur, rx).await {
                Ok(r) => r,
                Err(_) => {
                    self.shared.pending.lock().await.remove(&id);
                    return Err(ClientError::Cancelled);
                }
            },
            None => rx.await,
        };

        match result {
            Ok(Ok(value)) => Ok(serde_json::from_value(value)?),
            Ok(Err(e)) => Err(ClientError::Rpc(e)),
            Err(_) => Err(ClientError::Shutdown),
        }
    }

    /// Send a JSON-RPC notification (fire-and-forget).
    pub async fn notify<P: Serialize>(&self, method: &str, params: P) -> Result<(), ClientError> {
        let params_val = serde_json::to_value(&params)?;
        let params_any = if params_val.is_null() {
            None
        } else {
            Some(ahp_types::common::AnyValue::from(params_val))
        };
        let msg = JsonRpcMessage::Notification(JsonRpcNotification {
            jsonrpc: JsonRpcVersion::V2,
            method: method.into(),
            params: params_any,
        });
        self.shared
            .outbound
            .send(Outbound::Message(msg))
            .await
            .map_err(|_| ClientError::Shutdown)
    }

    /// Issue the `initialize` handshake.
    pub async fn initialize(
        &self,
        client_id: String,
        protocol_version: i64,
        initial_subscriptions: Vec<String>,
    ) -> Result<InitializeResult, ClientError> {
        let params = InitializeParams {
            protocol_version,
            client_id,
            initial_subscriptions: if initial_subscriptions.is_empty() {
                None
            } else {
                Some(initial_subscriptions)
            },
            locale: None,
        };
        self.request("initialize", params).await
    }

    /// Re-establish a dropped connection with `reconnect`.
    pub async fn reconnect(
        &self,
        client_id: String,
        last_seen_server_seq: i64,
        subscriptions: Vec<String>,
    ) -> Result<ReconnectResult, ClientError> {
        let params = ReconnectParams {
            client_id,
            last_seen_server_seq,
            subscriptions,
        };
        self.request("reconnect", params).await
    }

    /// Subscribe to a URI and obtain a handle that streams
    /// [`SubscriptionEvent`]s for that resource.
    pub async fn subscribe(
        &self,
        uri: String,
    ) -> Result<(SubscribeResult, SessionSubscription), ClientError> {
        let sub = self.attach_subscription(&uri).await;
        let result: SubscribeResult = self
            .request("subscribe", SubscribeParams { resource: uri })
            .await?;
        Ok((result, sub))
    }

    /// Attach a new local subscription handle for `uri` without sending a
    /// `subscribe` request — useful when the URI was included in
    /// `initialSubscriptions` during [`Client::initialize`].
    pub async fn attach_subscription(&self, uri: &str) -> SessionSubscription {
        let mut subs = self.shared.subscriptions.lock().await;
        let tx = subs
            .entry(uri.to_string())
            .or_insert_with(|| broadcast::channel(self.shared.config.subscription_buffer).0);
        SessionSubscription {
            rx: tx.subscribe(),
            uri: uri.to_string(),
        }
    }

    /// Send an `unsubscribe` notification and drop the local fan-out for
    /// this URI. In-flight [`SessionSubscription`] handles will see the
    /// channel close.
    pub async fn unsubscribe(&self, uri: String) -> Result<(), ClientError> {
        {
            let mut subs = self.shared.subscriptions.lock().await;
            subs.remove(&uri);
        }
        self.notify("unsubscribe", UnsubscribeParams { resource: uri })
            .await
    }

    /// Fire a write-ahead `dispatchAction` notification with a
    /// client-assigned sequence number.
    pub async fn dispatch(&self, action: StateAction) -> Result<DispatchHandle, ClientError> {
        let client_seq = self.shared.next_client_seq.fetch_add(1, Ordering::Relaxed) as i64;
        self.notify(
            "dispatchAction",
            DispatchActionParams { client_seq, action },
        )
        .await?;
        Ok(DispatchHandle { client_seq })
    }
}

async fn drive_transport<T: Transport>(
    mut transport: T,
    shared: Arc<Shared>,
    mut outbound: mpsc::Receiver<Outbound>,
) {
    loop {
        tokio::select! {
            outbound_msg = outbound.recv() => {
                match outbound_msg {
                    Some(Outbound::Message(msg)) => {
                        if let Ok(wire) = TransportMessage::encode(&msg) {
                            if let Err(err) = transport.send(wire).await {
                                tracing::warn!(?err, "transport send failed");
                                break;
                            }
                        }
                    }
                    Some(Outbound::Shutdown) | None => {
                        let _ = transport.close().await;
                        break;
                    }
                }
            }
            inbound = transport.recv() => {
                match inbound {
                    Ok(Some(wire)) => {
                        match wire.into_parsed() {
                            Ok(msg) => dispatch_inbound(&shared, msg).await,
                            Err(err) => tracing::warn!(?err, "malformed frame"),
                        }
                    }
                    Ok(None) => break,
                    Err(err) => {
                        tracing::warn!(?err, "transport recv error");
                        break;
                    }
                }
            }
        }
    }

    // Teardown: close everything so waiters see Shutdown.
    let mut pending = shared.pending.lock().await;
    for (_, tx) in pending.drain() {
        let _ = tx.send(Err(JsonRpcError {
            code: -32000,
            message: "transport closed".into(),
            data: None,
        }));
    }
    let mut subs = shared.subscriptions.lock().await;
    subs.clear();
}

async fn dispatch_inbound(shared: &Shared, msg: JsonRpcMessage) {
    match msg {
        JsonRpcMessage::SuccessResponse(r) => {
            if let Some(tx) = shared.pending.lock().await.remove(&r.id) {
                let _ = tx.send(Ok(r.result));
            }
        }
        JsonRpcMessage::ErrorResponse(r) => {
            if let Some(tx) = shared.pending.lock().await.remove(&r.id) {
                let _ = tx.send(Err(r.error));
            }
        }
        JsonRpcMessage::Notification(n) => {
            handle_notification(shared, n).await;
        }
        JsonRpcMessage::Request(r) => {
            tracing::debug!(method = %r.method, "ignoring unexpected server request");
        }
    }
}

async fn handle_notification(shared: &Shared, n: JsonRpcNotification) {
    let params_val: Value = n.params.unwrap_or(Value::Null);

    match n.method.as_str() {
        "action" => {
            if let Ok(envelope) = serde_json::from_value::<ActionNotificationParams>(params_val) {
                let resource = action_resource(&envelope.action);
                let subs = shared.subscriptions.lock().await;
                if let Some(resource) = resource {
                    if let Some(tx) = subs.get(&resource) {
                        let _ = tx.send(SubscriptionEvent::Action(envelope.clone()));
                    }
                }
            }
        }
        "notification" => {
            if let Ok(wrapped) = serde_json::from_value::<NotificationMethodParams>(params_val) {
                let subs = shared.subscriptions.lock().await;
                // Protocol notifications are cross-resource; fan them out
                // to every active subscription. Callers that care about a
                // specific notification can filter.
                for tx in subs.values() {
                    let _ = tx.send(SubscriptionEvent::Notification(
                        wrapped.notification.clone(),
                    ));
                }
            }
        }
        other => {
            tracing::debug!(method = %other, "unhandled notification");
        }
    }
}

/// Extract the resource URI an action is scoped to.
///
/// Rather than enumerating every variant (which drifts as new actions are
/// generated), we serialize the payload and inspect the `session` or
/// `terminal` field. Root actions have neither and fall back to the
/// well-known `root:/` URI.
fn action_resource(action: &StateAction) -> Option<String> {
    let val = serde_json::to_value(action).ok()?;
    if let Some(uri) = val.get("session").and_then(Value::as_str) {
        return Some(uri.to_string());
    }
    if let Some(uri) = val.get("terminal").and_then(Value::as_str) {
        return Some(uri.to_string());
    }
    Some("root:/".into())
}