use std::{
    pin::Pin,
    sync::{atomic::{AtomicU32, Ordering}, Arc},
    task::{Context, Poll},
    time::Duration,
};
use futures_util::Stream;
use tokio::{
    net::TcpStream,
    sync::{oneshot, mpsc},
    time::timeout,
};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use crate::{
    error::{CdpError, Result},
    handler::Handler,
    types::{CallId, Message, MethodCall, MethodId, SessionId},
};

/// Specifies the configuration for a `Connection`.
#[derive(Debug, Clone)]
pub struct ConnectionConfig {
    /// The default timeout for sending a command to the browser.
    pub default_timeout: Duration,
    /// The maximum number of commands that can be buffered before new commands
    /// will be rejected. This is the capacity of the command channel.
    pub max_pending_commands: usize,
    /// The maximum number of events that can be buffered. This is the capacity
    /// of the event channel.
    pub max_buffered_events: usize,
}

impl Default for ConnectionConfig {
    fn default() -> Self {
        Self {
            default_timeout: Duration::from_secs(30),
            max_pending_commands: 1000,
            max_buffered_events: 10000,
        }
    }
}

/// Represents an active WebSocket connection to a Chrome DevTools Protocol server.
///
/// `Connection` is the primary entry point for interacting with the browser.
/// It allows sending commands and receiving events.
///
/// It manages the connection state, command dispatch, and event stream behind
/// the scenes using a dedicated background task.
///
/// # Buffering and Backpressure
///
/// The connection uses bounded channels for both commands and events to manage memory
/// usage and apply backpressure.
/// - The command channel is limited by `max_pending_commands` in `ConnectionConfig`.
/// - The event channel is limited by `max_buffered_events` in `ConnectionConfig`.
///
/// If the command channel is full, `call_method` and related methods will return
/// an error immediately, preventing the command queue from growing indefinitely.
#[derive(Debug)]
pub struct Connection {
    /// A channel for sending `MethodCall` instances to the background handler task.
    command_tx: mpsc::Sender<(MethodCall, oneshot::Sender<Result<serde_json::Value>>)>,
    /// A channel for receiving `Message` events from the background handler task.
    event_rx: mpsc::Receiver<Message>,
    /// An atomic counter to generate unique IDs for method calls.
    next_id: Arc<AtomicU32>,
    /// The configuration for this connection.
    config: ConnectionConfig,
    /// A channel to signal the background handler task to shut down.
    shutdown_tx: mpsc::UnboundedSender<()>,
}

impl Connection {
    /// Establishes a new `Connection` over an existing WebSocket stream.
    ///
    /// This method creates a `Connection` with the default `ConnectionConfig`
    /// and spawns a background task to manage the WebSocket communication.
    ///
    /// # Arguments
    ///
    /// * `ws`: The `WebSocketStream` connected to the Chrome DevTools Protocol endpoint.
    ///
    /// # Returns
    ///
    /// A new `Connection` instance.
    pub fn new(ws: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
        Self::with_config(ws, ConnectionConfig::default())
    }
    
    /// Establishes a new `Connection` with a specific default command timeout.
    ///
    /// All other configuration options will use their default values.
    ///
    /// # Arguments
    ///
    /// * `ws`: The `WebSocketStream` connected to the Chrome DevTools Protocol endpoint.
    /// * `default_timeout`: The default `Duration` to wait for a command response.
    ///
    /// # Returns
    ///
    /// A new `Connection` instance.
    pub fn with_timeout(
        ws: WebSocketStream<MaybeTlsStream<TcpStream>>,
        default_timeout: Duration,
    ) -> Self {
        let config = ConnectionConfig {
            default_timeout,
            ..Default::default()
        };
        Self::with_config(ws, config)
    }
    
    /// Establishes a new `Connection` with a custom configuration.
    ///
    /// This is the most flexible way to create a `Connection`, allowing you to
    /// control timeouts and buffer sizes.
    ///
    /// # Arguments
    ///
    /// * `ws`: The `WebSocketStream` connected to the Chrome DevTools Protocol endpoint.
    /// * `config`: The `ConnectionConfig` to use for this connection.
    ///
    /// # Returns
    ///
    /// A new `Connection` instance.
    pub fn with_config(
        ws: WebSocketStream<MaybeTlsStream<TcpStream>>,
        config: ConnectionConfig,
    ) -> Self {
        // Create bounded channels to manage backpressure and prevent unbounded memory usage.
        let (command_tx, command_rx) = mpsc::channel(config.max_pending_commands);
        let (event_tx, event_rx) = mpsc::channel(config.max_buffered_events);
        let (shutdown_tx, shutdown_rx) = mpsc::unbounded_channel();
        let next_id = Arc::new(AtomicU32::new(1));
        
        // Spawn the background handler task to manage the WebSocket connection.
        let handler = Handler::new(
            command_rx,
            event_tx,
            shutdown_rx,
            ws,
        );
        tokio::spawn(handler.run());
        
        Self {
            command_tx,
            event_rx,
            next_id,
            config,
            shutdown_tx,
        }
    }
    
    /// Atomically generates a new, unique `CallId` for a method call.
    fn next_call_id(&self) -> CallId {
        CallId::new(self.next_id.fetch_add(1, Ordering::Relaxed))
    }
    
    /// Sends a method call to the browser and waits for the result.
    ///
    /// This uses the `default_timeout` from the `ConnectionConfig`.
    ///
    /// # Arguments
    ///
    /// * `method`: The name of the method to call, e.g., `"Target.getTargets"`.
    /// * `params`: The parameters for the method call, as a `serde_json::Value`.
    ///
    /// # Returns
    ///
    /// A `Result` containing the successful JSON response from the browser,
    /// or a `CdpError` if the command fails, times out, or the connection is closed.
    pub async fn call_method(
        &mut self,
        method: impl Into<MethodId>,
        params: serde_json::Value,
    ) -> Result<serde_json::Value> {
        self.call_method_with_timeout(method, params, self.config.default_timeout)
            .await
    }
    
    /// Sends a method call to the browser with a specific timeout.
    ///
    /// # Arguments
    ///
    /// * `method`: The name of the method to call.
    /// * `params`: The parameters for the method call.
    /// * `timeout_duration`: The maximum `Duration` to wait for a response.
    ///
    /// # Returns
    ///
    /// A `Result` containing the successful JSON response from the browser,
    /// or a `CdpError` if the command fails, times out, or the connection is closed.
    pub async fn call_method_with_timeout(
        &mut self,
        method: impl Into<MethodId>,
        params: serde_json::Value,
        timeout_duration: Duration,
    ) -> Result<serde_json::Value> {
        let call_id = self.next_call_id();
        let call = MethodCall::new(call_id, method, params);
        self.dispatch_command(call, timeout_duration).await
    }
    
    /// Sends a method call to a specific session and waits for the result.
    ///
    /// This is used to interact with a specific target, like a page or a worker,
    /// that has an active session. This uses the `default_timeout` from the `ConnectionConfig`.
    ///
    /// # Arguments
    ///
    /// * `method`: The name of the method to call.
    /// * `session_id`: The ID of the session to send the command to.
    /// * `params`: The parameters for the method call.
    ///
    /// # Returns
    ///
    /// A `Result` containing the successful JSON response from the browser,
    /// or a `CdpError` if the command fails, times out, or the connection is closed.
    pub async fn call_method_in_session(
        &mut self,
        method: impl Into<MethodId>,
        session_id: impl Into<SessionId>,
        params: serde_json::Value,
    ) -> Result<serde_json::Value> {
        self.call_method_in_session_with_timeout(
            method, 
            session_id, 
            params, 
            self.config.default_timeout
        ).await
    }
    
    /// Sends a method call to a specific session with a specific timeout.
    ///
    /// # Arguments
    ///
    /// * `method`: The name of the method to call.
    /// * `session_id`: The ID of the session to send the command to.
    /// * `params`: The parameters for the method call.
    /// * `timeout_duration`: The maximum `Duration` to wait for a response.
    ///
    /// # Returns
    ///
    /// A `Result` containing the successful JSON response from the browser,
    /// or a `CdpError` if the command fails, times out, or the connection is closed.
    pub async fn call_method_in_session_with_timeout(
        &mut self,
        method: impl Into<MethodId>,
        session_id: impl Into<SessionId>,
        params: serde_json::Value,
        timeout_duration: Duration,
    ) -> Result<serde_json::Value> {
        let call_id = self.next_call_id();
        let call = MethodCall::with_session(call_id, method, session_id, params);
        self.dispatch_command(call, timeout_duration).await
    }
    
    /// Helper function to dispatch a command and wait for its response.
    async fn dispatch_command(
        &mut self,
        call: MethodCall,
        timeout_duration: Duration,
    ) -> Result<serde_json::Value> {
        let (response_tx, response_rx) = oneshot::channel();
        
        // Try to send the command to the handler, returning an error if the channel is full.
        self.command_tx
            .try_send((call, response_tx))
            .map_err(|e| match e {
                mpsc::error::TrySendError::Full(_) => CdpError::ProcessError(
                    "Command queue is full, too many pending commands".to_string(),
                ),
                mpsc::error::TrySendError::Closed(_) => CdpError::ConnectionClosed,
            })?;
        
        // Wait for the response from the handler.
        match timeout(timeout_duration, response_rx).await {
            Ok(Ok(result)) => result,
            Ok(Err(_)) => Err(CdpError::InvalidResponse("Response channel closed".to_string())),
            Err(_) => Err(CdpError::TimeoutError {
                timeout_ms: timeout_duration.as_millis() as u64,
            }),
        }
    }
    
    /// Returns a reference to the `ConnectionConfig` used by this connection.
    pub fn config(&self) -> &ConnectionConfig {
        &self.config
    }

    /// Signals the connection to gracefully shut down.
    ///
    /// This sends a shutdown signal to the background handler task. The task will
    /// attempt to close the WebSocket connection gracefully. Any pending commands
    /// may fail.
    ///
    /// The `Connection` will become unusable after this method is called.
    /// Note that `shutdown` is also called automatically when the `Connection` is dropped.
    pub fn shutdown(&mut self) -> Result<()> {
        self.shutdown_tx
            .send(())
            .map_err(|_| CdpError::ConnectionClosed)?;
        Ok(())
    }
}

impl Stream for Connection {
    type Item = Result<Message>;
    
    /// Polls for the next incoming event from the browser.
    ///
    /// This method is part of the `Stream` implementation and allows the `Connection`
    /// to be used as an asynchronous stream of `Message` events.
    ///
    /// # Returns
    ///
    /// * `Poll::Ready(Some(Ok(message)))`: An event was successfully received.
    /// * `Poll::Ready(Some(Err(e)))`: An error occurred, typically `CdpError::ConnectionClosed`.
    /// * `Poll::Ready(None)`: The stream is closed and will not produce more items.
    /// * `Poll::Pending`: No event is currently available. The waker will be notified
    ///   when a new event arrives.
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        
        match this.event_rx.poll_recv(cx) {
            Poll::Ready(Some(event)) => Poll::Ready(Some(Ok(event))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl Drop for Connection {
    /// Ensures that the connection is shut down when it goes out of scope.
    ///
    /// This sends a shutdown signal to the background task, allowing for a graceful
    /// cleanup of the WebSocket connection and associated resources.
    fn drop(&mut self) {
        // Attempt to send a shutdown signal.
        // We ignore the result because the receiver might have already been dropped
        // if the connection was explicitly shut down or the handler task terminated.
        let _ = self.shutdown_tx.send(());
    }
}

Analyze this code