cdp-use-rs 0.1.1

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use futures_util::{SinkExt, StreamExt};
use serde_json::Value;
use tokio::net::TcpStream;
use tokio::sync::{oneshot, Mutex as AsyncMutex};
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async_with_config, MaybeTlsStream, WebSocketStream};

use crate::CdpError;

/// Default timeout for CDP commands (30 seconds).
const DEFAULT_COMMAND_TIMEOUT: Duration = Duration::from_secs(30);

/// Default maximum WebSocket message size (100 MiB, matching Python cdp-use).
const DEFAULT_MAX_MESSAGE_SIZE: usize = 100 * 1024 * 1024;

/// Configuration for a CDP client connection.
///
/// Use `Default::default()` for sensible defaults matching the Python cdp-use client.
#[derive(Debug, Clone)]
pub struct CdpClientConfig {
    /// Maximum WebSocket message size in bytes. Default: 100 MiB.
    pub max_message_size: Option<usize>,
    /// Maximum WebSocket frame size in bytes. Default: tungstenite default (16 MiB).
    pub max_frame_size: Option<usize>,
    /// Additional HTTP headers to send during the WebSocket handshake.
    pub additional_headers: HashMap<String, String>,
    /// Timeout for CDP commands. Default: 30 seconds.
    pub command_timeout: Duration,
}

impl Default for CdpClientConfig {
    fn default() -> Self {
        Self {
            max_message_size: Some(DEFAULT_MAX_MESSAGE_SIZE),
            max_frame_size: None, // use tungstenite default (16 MiB)
            additional_headers: HashMap::new(),
            command_timeout: DEFAULT_COMMAND_TIMEOUT,
        }
    }
}

type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
type WsSink = futures_util::stream::SplitSink<WsStream, Message>;
type WsSource = futures_util::stream::SplitStream<WsStream>;
type PendingRequests = HashMap<u64, oneshot::Sender<Result<Value, CdpError>>>;

/// Type-erased async event handler.
pub type EventHandler = Arc<
    dyn Fn(Value, Option<String>) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync,
>;

/// Central registry for CDP event callbacks.
///
/// One handler per method (replacement semantics, matching the Python version).
pub struct EventRegistry {
    handlers: std::sync::Mutex<HashMap<String, EventHandler>>,
}

impl EventRegistry {
    pub fn new() -> Self {
        Self {
            handlers: std::sync::Mutex::new(HashMap::new()),
        }
    }

    /// Register a handler for a CDP event method. Replaces any existing handler.
    pub fn register(&self, method: &str, handler: EventHandler) {
        self.handlers
            .lock()
            .unwrap()
            .insert(method.to_string(), handler);
    }

    /// Remove the handler for a CDP event method.
    pub fn unregister(&self, method: &str) {
        self.handlers.lock().unwrap().remove(method);
    }

    /// Dispatch an event to its registered handler. Returns true if handled.
    ///
    /// The handler is cloned and the lock is dropped before awaiting, so
    /// handlers may safely call `register`/`unregister` without deadlocking.
    pub async fn handle_event(
        &self,
        method: &str,
        params: Value,
        session_id: Option<String>,
    ) -> bool {
        let handler = {
            let handlers = self.handlers.lock().unwrap();
            handlers.get(method).cloned()
        };

        if let Some(handler) = handler {
            handler(params, session_id).await;
            true
        } else {
            false
        }
    }

    /// Remove all registered handlers.
    pub fn clear(&self) {
        self.handlers.lock().unwrap().clear();
    }
}

impl Default for EventRegistry {
    fn default() -> Self {
        Self::new()
    }
}

/// Chrome DevTools Protocol client.
///
/// Connects to a browser via WebSocket and provides access to CDP
/// commands and events.
///
/// ```no_run
/// # async fn example() -> Result<(), cdp_use::CdpError> {
/// let cdp = cdp_use::CdpClient::connect("ws://127.0.0.1:9222/devtools/browser/...").await?;
/// let result = cdp.send_raw("Target.getTargets", serde_json::json!({}), None).await?;
/// cdp.close().await?;
/// # Ok(())
/// # }
/// ```
#[derive(Clone)]
pub struct CdpClient {
    inner: Arc<ClientInner>,
}

struct ClientInner {
    sink: AsyncMutex<WsSink>,
    next_id: AtomicU64,
    pending: Arc<AsyncMutex<PendingRequests>>,
    event_registry: Arc<EventRegistry>,
    closed: AtomicBool,
    command_timeout: Duration,
    message_loop_handle: std::sync::Mutex<Option<JoinHandle<()>>>,
}

impl Drop for ClientInner {
    fn drop(&mut self) {
        if let Some(handle) = self.message_loop_handle.get_mut().unwrap().take() {
            handle.abort();
        }
    }
}

impl CdpClient {
    /// Connect to a CDP endpoint via WebSocket with default configuration.
    pub async fn connect(url: &str) -> Result<Self, CdpError> {
        Self::connect_with_config(url, CdpClientConfig::default()).await
    }

    /// Connect to a CDP endpoint via WebSocket with custom configuration.
    ///
    /// ```no_run
    /// # async fn example() -> Result<(), cdp_use::CdpError> {
    /// use cdp_use::{CdpClient, CdpClientConfig};
    /// use std::time::Duration;
    ///
    /// let config = CdpClientConfig {
    ///     command_timeout: Duration::from_secs(60),
    ///     ..Default::default()
    /// };
    /// let cdp = CdpClient::connect_with_config("ws://127.0.0.1:9222/devtools/browser/...", config).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn connect_with_config(
        url: &str,
        config: CdpClientConfig,
    ) -> Result<Self, CdpError> {
        let mut request = url.into_client_request()?;

        // Add custom headers to the WebSocket handshake request
        for (key, value) in &config.additional_headers {
            request.headers_mut().insert(
                key.parse::<tokio_tungstenite::tungstenite::http::HeaderName>()
                    .map_err(|e| CdpError::Protocol {
                        code: -1,
                        message: format!("Invalid header name '{key}': {e}"),
                        data: None,
                    })?,
                value
                    .parse()
                    .map_err(|e| CdpError::Protocol {
                        code: -1,
                        message: format!("Invalid header value for '{key}': {e}"),
                        data: None,
                    })?,
            );
        }

        let mut ws_config = WebSocketConfig::default();
        ws_config.max_message_size = config.max_message_size;
        ws_config.max_frame_size = config.max_frame_size;

        let (ws_stream, _) =
            connect_async_with_config(request, Some(ws_config), false).await?;
        let (sink, stream) = ws_stream.split();

        let pending = Arc::new(AsyncMutex::new(HashMap::new()));
        let event_registry = Arc::new(EventRegistry::new());
        let closed = Arc::new(AtomicBool::new(false));

        let handle = tokio::spawn({
            let pending = pending.clone();
            let registry = event_registry.clone();
            let closed = closed.clone();
            async move {
                message_loop(stream, pending, registry, closed).await;
            }
        });

        Ok(Self {
            inner: Arc::new(ClientInner {
                sink: AsyncMutex::new(sink),
                next_id: AtomicU64::new(0),
                pending,
                event_registry,
                closed: AtomicBool::new(false),
                command_timeout: config.command_timeout,
                message_loop_handle: std::sync::Mutex::new(Some(handle)),
            }),
        })
    }

    /// Send a raw CDP command and await the response.
    ///
    /// Returns `CdpError::ConnectionClosed` if the connection is already closed,
    /// or `CdpError::Timeout` if the browser does not respond within the timeout.
    pub async fn send_raw(
        &self,
        method: &str,
        params: Value,
        session_id: Option<&str>,
    ) -> Result<Value, CdpError> {
        if self.inner.closed.load(Ordering::Acquire) {
            return Err(CdpError::ConnectionClosed);
        }

        let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed) + 1;

        let (tx, rx) = oneshot::channel();
        self.inner.pending.lock().await.insert(id, tx);

        let mut msg = serde_json::json!({
            "id": id,
            "method": method,
            "params": params,
        });
        if let Some(sid) = session_id {
            msg["sessionId"] = Value::String(sid.to_string());
        }

        let send_result = self
            .inner
            .sink
            .lock()
            .await
            .send(Message::Text(msg.to_string().into()))
            .await;

        if let Err(e) = send_result {
            // Clean up the pending entry if the WebSocket send failed
            self.inner.pending.lock().await.remove(&id);
            return Err(e.into());
        }

        // Await response with timeout
        match tokio::time::timeout(self.inner.command_timeout, rx).await {
            Ok(Ok(result)) => result,
            Ok(Err(_)) => {
                // Sender dropped (connection closed)
                Err(CdpError::ConnectionClosed)
            }
            Err(_elapsed) => {
                // Timeout — clean up the pending entry
                self.inner.pending.lock().await.remove(&id);
                Err(CdpError::Timeout)
            }
        }
    }

    /// Emit a synthetic event through the event registry.
    ///
    /// Useful for custom domains where events are produced by application
    /// code rather than the browser.
    pub async fn emit_event(
        &self,
        method: &str,
        params: Value,
        session_id: Option<&str>,
    ) -> bool {
        self.inner
            .event_registry
            .handle_event(method, params, session_id.map(String::from))
            .await
    }

    /// Get a reference to the event registry.
    // Used by generated code in generated.rs
    pub(crate) fn event_registry(&self) -> &Arc<EventRegistry> {
        &self.inner.event_registry
    }

    /// Gracefully close the WebSocket connection.
    pub async fn close(&self) -> Result<(), CdpError> {
        // Mark as closed first so new send_raw calls fail fast
        self.inner.closed.store(true, Ordering::Release);

        // Fail all pending requests before aborting the message loop,
        // so in-flight callers get ConnectionClosed instead of hanging.
        {
            let mut pending = self.inner.pending.lock().await;
            for (_, tx) in pending.drain() {
                let _ = tx.send(Err(CdpError::ConnectionClosed));
            }
        }

        // Abort the message handler
        if let Some(handle) = self.inner.message_loop_handle.lock().unwrap().take() {
            handle.abort();
            let _ = handle.await;
        }

        // Send close frame
        self.inner.sink.lock().await.close().await?;
        Ok(())
    }
}

async fn message_loop(
    mut stream: WsSource,
    pending: Arc<AsyncMutex<PendingRequests>>,
    event_registry: Arc<EventRegistry>,
    closed: Arc<AtomicBool>,
) {
    while let Some(msg_result) = stream.next().await {
        match msg_result {
            Ok(Message::Text(text)) => {
                let data: Value = match serde_json::from_str(&text) {
                    Ok(v) => v,
                    Err(e) => {
                        tracing::warn!("Failed to parse CDP message: {e}");
                        continue;
                    }
                };

                if let Some(id) = data.get("id").and_then(|v| v.as_u64()) {
                    // Response message — match to pending request
                    let mut pending = pending.lock().await;
                    if let Some(tx) = pending.remove(&id) {
                        let result = if let Some(error) = data.get("error") {
                            let code = error.get("code").and_then(|v| v.as_i64()).unwrap_or(0);
                            let message = error
                                .get("message")
                                .and_then(|v| v.as_str())
                                .unwrap_or("Unknown error")
                                .to_string();
                            let err_data = error.get("data").map(|v| v.to_string());
                            Err(CdpError::Protocol {
                                code,
                                message,
                                data: err_data,
                            })
                        } else {
                            Ok(data
                                .get("result")
                                .cloned()
                                .unwrap_or(Value::Object(Default::default())))
                        };
                        let _ = tx.send(result);
                    }
                } else if let Some(method) = data.get("method").and_then(|v| v.as_str()) {
                    // Event message — spawn handler to avoid blocking the message loop
                    let params = data.get("params").cloned().unwrap_or_default();
                    let session_id = data
                        .get("sessionId")
                        .and_then(|v| v.as_str())
                        .map(String::from);
                    let registry = event_registry.clone();
                    let method = method.to_string();
                    tokio::spawn(async move {
                        registry.handle_event(&method, params, session_id).await;
                    });
                }
            }
            Ok(Message::Close(_)) | Err(_) => {
                // Connection closed or error — mark closed and fail all pending requests
                closed.store(true, Ordering::Release);
                let mut pending = pending.lock().await;
                for (_, tx) in pending.drain() {
                    let _ = tx.send(Err(CdpError::ConnectionClosed));
                }
                break;
            }
            _ => {} // Ping/pong handled automatically by tokio-tungstenite
        }
    }
}