adk_browser/

session.rs

//! Browser session management wrapping thirtyfour WebDriver.

use crate::config::{BrowserConfig, BrowserType};
use crate::escape::escape_js_string;
use adk_core::{AdkError, Result};
use std::sync::Arc;
use std::time::Duration;
use thirtyfour::common::print::{PrintOrientation, PrintParameters};
use thirtyfour::prelude::*;
use tokio::sync::RwLock;

/// State information about an element.
#[derive(Debug, Clone)]
pub struct ElementState {
    /// Whether the element is displayed (visible).
    pub is_displayed: bool,
    /// Whether the element is enabled (not disabled).
    pub is_enabled: bool,
    /// Whether the element is selected (for checkboxes, radio buttons, options).
    pub is_selected: bool,
    /// Whether the element is clickable (displayed and enabled).
    pub is_clickable: bool,
}

/// A browser session that wraps thirtyfour's WebDriver.
///
/// This is the core abstraction for browser automation in ADK.
/// It can be shared across multiple tools via `Arc<BrowserSession>`.
pub struct BrowserSession {
    driver: RwLock<Option<WebDriver>>,
    config: BrowserConfig,
}

impl BrowserSession {
    /// Create a new browser session with the given configuration.
    ///
    /// Note: This does not start the browser immediately.
    /// Call `start()` to initialize the WebDriver connection.
    pub fn new(config: BrowserConfig) -> Self {
        Self { driver: RwLock::new(None), config }
    }

    /// Create a browser session with default configuration.
    pub fn with_defaults() -> Self {
        Self::new(BrowserConfig::default())
    }

    /// Start the browser session by connecting to WebDriver.
    pub async fn start(&self) -> Result<()> {
        let mut driver_guard = self.driver.write().await;

        if driver_guard.is_some() {
            return Ok(()); // Already started
        }

        let caps = self.build_capabilities()?;
        let driver = WebDriver::new(&self.config.webdriver_url, caps)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to start browser: {}", e)))?;

        // Set timeouts
        driver
            .set_page_load_timeout(Duration::from_secs(self.config.page_load_timeout_secs))
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to set page load timeout: {}", e)))?;

        driver
            .set_script_timeout(Duration::from_secs(self.config.script_timeout_secs))
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to set script timeout: {}", e)))?;

        driver
            .set_implicit_wait_timeout(Duration::from_secs(self.config.implicit_wait_secs))
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to set implicit wait: {}", e)))?;

        // Set viewport size
        driver
            .set_window_rect(0, 0, self.config.viewport_width, self.config.viewport_height)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to set viewport: {}", e)))?;

        *driver_guard = Some(driver);
        Ok(())
    }

    /// Stop the browser session.
    pub async fn stop(&self) -> Result<()> {
        let mut driver_guard = self.driver.write().await;

        if let Some(driver) = driver_guard.take() {
            driver
                .quit()
                .await
                .map_err(|e| AdkError::Tool(format!("Failed to quit browser: {}", e)))?;
        }

        Ok(())
    }

    /// Check if the session is active by verifying the WebDriver connection is alive.
    pub async fn is_active(&self) -> bool {
        let driver_guard = self.driver.read().await;
        if let Some(ref driver) = *driver_guard {
            // Actually ping the session to verify it's alive
            driver.title().await.is_ok()
        } else {
            false
        }
    }

    /// Ensure the browser session is started, reconnecting if the session died.
    ///
    /// Call this instead of checking `is_active()` manually. If the session
    /// exists but is stale (WebDriver died), it will be recreated transparently.
    pub async fn ensure_started(&self) -> Result<()> {
        {
            let driver_guard = self.driver.read().await;
            if let Some(ref driver) = *driver_guard {
                // Ping the session — if it responds, we're good
                if driver.title().await.is_ok() {
                    return Ok(());
                }
            }
        }
        // Session is dead or missing — (re)create it
        // First clear the stale driver if any
        {
            let mut driver_guard = self.driver.write().await;
            *driver_guard = None;
        }
        self.start().await
    }

    /// Internal: get a live WebDriver handle, auto-starting if needed.
    /// All public methods that access the WebDriver MUST go through this.
    ///
    /// Returns a cloned WebDriver handle (cheap — WebDriver is Arc-based internally).
    /// The RwLock is NOT held across the caller's WebDriver operations.
    async fn live_driver(&self) -> Result<WebDriver> {
        self.ensure_started().await?;

        let guard = self.driver.read().await;
        guard.clone().ok_or_else(|| {
            AdkError::Tool(
                "Failed to start browser session: driver is None after ensure_started()".into(),
            )
        })
    }

    /// Get a snapshot of the current page context (url, title, truncated body text).
    ///
    /// Useful for returning page state after interaction tools.
    pub async fn page_context(&self) -> Result<serde_json::Value> {
        let url = self.current_url().await.unwrap_or_default();
        let title = self.title().await.unwrap_or_default();

        // Get truncated visible text (first 2000 chars of body)
        let page_text = self
            .execute_script(
                "return (document.body && document.body.innerText || '').substring(0, 2000);",
            )
            .await
            .ok()
            .and_then(|v| v.as_str().map(|s| s.to_string()))
            .unwrap_or_default();

        Ok(serde_json::json!({
            "url": url,
            "title": title,
            "page_text": page_text
        }))
    }

    /// Get the configuration.
    pub fn config(&self) -> &BrowserConfig {
        &self.config
    }

    /// Navigate to a URL.
    pub async fn navigate(&self, url: &str) -> Result<()> {
        let driver = self.live_driver().await?;

        driver.goto(url).await.map_err(|e| AdkError::Tool(format!("Navigation failed: {}", e)))?;

        Ok(())
    }

    /// Get the current URL.
    pub async fn current_url(&self) -> Result<String> {
        let driver = self.live_driver().await?;

        driver
            .current_url()
            .await
            .map(|u| u.to_string())
            .map_err(|e| AdkError::Tool(format!("Failed to get URL: {}", e)))
    }

    /// Get the page title.
    pub async fn title(&self) -> Result<String> {
        let driver = self.live_driver().await?;

        driver.title().await.map_err(|e| AdkError::Tool(format!("Failed to get title: {}", e)))
    }

    /// Find an element by CSS selector.
    pub async fn find_element(&self, selector: &str) -> Result<WebElement> {
        let driver = self.live_driver().await?;

        driver
            .find(By::Css(selector))
            .await
            .map_err(|e| AdkError::Tool(format!("Element not found '{}': {}", selector, e)))
    }

    /// Find multiple elements by CSS selector.
    pub async fn find_elements(&self, selector: &str) -> Result<Vec<WebElement>> {
        let driver = self.live_driver().await?;

        driver
            .find_all(By::Css(selector))
            .await
            .map_err(|e| AdkError::Tool(format!("Elements query failed '{}': {}", selector, e)))
    }

    /// Find element by XPath.
    pub async fn find_by_xpath(&self, xpath: &str) -> Result<WebElement> {
        let driver = self.live_driver().await?;

        driver
            .find(By::XPath(xpath))
            .await
            .map_err(|e| AdkError::Tool(format!("XPath not found '{}': {}", xpath, e)))
    }

    /// Click an element by selector.
    pub async fn click(&self, selector: &str) -> Result<()> {
        let element = self.find_element(selector).await?;
        element
            .click()
            .await
            .map_err(|e| AdkError::Tool(format!("Click failed on '{}': {}", selector, e)))
    }

    /// Type text into an element.
    pub async fn type_text(&self, selector: &str, text: &str) -> Result<()> {
        let element = self.find_element(selector).await?;
        element
            .send_keys(text)
            .await
            .map_err(|e| AdkError::Tool(format!("Type failed on '{}': {}", selector, e)))
    }

    /// Clear an input field.
    pub async fn clear(&self, selector: &str) -> Result<()> {
        let element = self.find_element(selector).await?;
        element
            .clear()
            .await
            .map_err(|e| AdkError::Tool(format!("Clear failed on '{}': {}", selector, e)))
    }

    /// Get text content of an element.
    pub async fn get_text(&self, selector: &str) -> Result<String> {
        let element = self.find_element(selector).await?;
        element
            .text()
            .await
            .map_err(|e| AdkError::Tool(format!("Get text failed on '{}': {}", selector, e)))
    }

    /// Get an attribute value.
    pub async fn get_attribute(&self, selector: &str, attribute: &str) -> Result<Option<String>> {
        let element = self.find_element(selector).await?;
        element
            .attr(attribute)
            .await
            .map_err(|e| AdkError::Tool(format!("Get attribute failed: {}", e)))
    }

    /// Take a screenshot (returns base64-encoded PNG).
    pub async fn screenshot(&self) -> Result<String> {
        let driver = self.live_driver().await?;

        let screenshot = driver
            .screenshot_as_png_base64()
            .await
            .map_err(|e| AdkError::Tool(format!("Screenshot failed: {}", e)))?;

        Ok(screenshot)
    }

    /// Take a screenshot of a specific element.
    pub async fn screenshot_element(&self, selector: &str) -> Result<String> {
        let element = self.find_element(selector).await?;
        let screenshot = element
            .screenshot_as_png_base64()
            .await
            .map_err(|e| AdkError::Tool(format!("Element screenshot failed: {}", e)))?;

        Ok(screenshot)
    }

    /// Execute JavaScript and return result.
    pub async fn execute_script(&self, script: &str) -> Result<serde_json::Value> {
        let driver = self.live_driver().await?;

        let result = driver
            .execute(script, vec![])
            .await
            .map_err(|e| AdkError::Tool(format!("Script execution failed: {}", e)))?;

        // thirtyfour's ScriptRet provides .json() which returns &Value directly
        Ok(result.json().clone())
    }

    /// Execute async JavaScript and return result.
    pub async fn execute_async_script(&self, script: &str) -> Result<serde_json::Value> {
        let driver = self.live_driver().await?;

        let result = driver
            .execute_async(script, vec![])
            .await
            .map_err(|e| AdkError::Tool(format!("Async script failed: {}", e)))?;

        Ok(result.json().clone())
    }

    /// Wait for an element to be present.
    pub async fn wait_for_element(&self, selector: &str, timeout_secs: u64) -> Result<WebElement> {
        let driver = self.live_driver().await?;

        driver
            .query(By::Css(selector))
            .wait(Duration::from_secs(timeout_secs), Duration::from_millis(100))
            .first()
            .await
            .map_err(|e| {
                AdkError::Tool(format!(
                    "Timeout waiting for '{}' after {}s: {}",
                    selector, timeout_secs, e
                ))
            })
    }

    /// Wait for an element to be clickable.
    pub async fn wait_for_clickable(
        &self,
        selector: &str,
        timeout_secs: u64,
    ) -> Result<WebElement> {
        let driver = self.live_driver().await?;

        driver
            .query(By::Css(selector))
            .wait(Duration::from_secs(timeout_secs), Duration::from_millis(100))
            .and_clickable()
            .first()
            .await
            .map_err(|e| {
                AdkError::Tool(format!("Timeout waiting for clickable '{}': {}", selector, e))
            })
    }

    /// Get page source HTML.
    pub async fn page_source(&self) -> Result<String> {
        let driver = self.live_driver().await?;

        driver
            .source()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to get page source: {}", e)))
    }

    /// Go back in history.
    pub async fn back(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver.back().await.map_err(|e| AdkError::Tool(format!("Back navigation failed: {}", e)))
    }

    /// Go forward in history.
    pub async fn forward(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .forward()
            .await
            .map_err(|e| AdkError::Tool(format!("Forward navigation failed: {}", e)))
    }

    /// Refresh the current page.
    pub async fn refresh(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver.refresh().await.map_err(|e| AdkError::Tool(format!("Refresh failed: {}", e)))
    }

    // =========================================================================
    // Cookie Management
    // =========================================================================

    /// Get all cookies.
    pub async fn get_all_cookies(&self) -> Result<Vec<serde_json::Value>> {
        let driver = self.live_driver().await?;

        let cookies = driver
            .get_all_cookies()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to get cookies: {}", e)))?;

        Ok(cookies
            .into_iter()
            .map(|c| {
                serde_json::json!({
                    "name": c.name,
                    "value": c.value,
                    "domain": c.domain,
                    "path": c.path,
                    "secure": c.secure,
                })
            })
            .collect())
    }

    /// Get a cookie by name.
    pub async fn get_cookie(&self, name: &str) -> Result<serde_json::Value> {
        let driver = self.live_driver().await?;

        let cookie = driver
            .get_named_cookie(name)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to get cookie '{}': {}", name, e)))?;

        Ok(serde_json::json!({
            "name": cookie.name,
            "value": cookie.value,
            "domain": cookie.domain,
            "path": cookie.path,
            "secure": cookie.secure,
        }))
    }

    /// Add a cookie.
    #[allow(clippy::too_many_arguments)]
    pub async fn add_cookie(
        &self,
        name: &str,
        value: &str,
        domain: Option<&str>,
        path: Option<&str>,
        secure: Option<bool>,
        expiry: Option<i64>,
    ) -> Result<()> {
        let driver = self.live_driver().await?;

        let mut cookie = thirtyfour::Cookie::new(name, value);
        if let Some(d) = domain {
            cookie.set_domain(d);
        }
        if let Some(p) = path {
            cookie.set_path(p);
        }
        if let Some(s) = secure {
            cookie.set_secure(s);
        }
        if let Some(exp) = expiry {
            // thirtyfour Cookie expiry expects seconds since epoch
            cookie.set_expiry(exp);
        }

        driver
            .add_cookie(cookie)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to add cookie: {e}")))
    }

    /// Delete a cookie.
    pub async fn delete_cookie(&self, name: &str) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .delete_cookie(name)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to delete cookie: {}", e)))
    }

    /// Delete all cookies.
    pub async fn delete_all_cookies(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .delete_all_cookies()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to delete all cookies: {}", e)))
    }

    // =========================================================================
    // Window Management
    // =========================================================================

    /// List all windows/tabs.
    pub async fn list_windows(&self) -> Result<(Vec<String>, String)> {
        let driver = self.live_driver().await?;

        let windows = driver
            .windows()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to get windows: {}", e)))?;

        let current = driver
            .window()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to get current window: {}", e)))?;

        Ok((windows.into_iter().map(|w| w.to_string()).collect(), current.to_string()))
    }

    /// Open a new tab.
    pub async fn new_tab(&self) -> Result<String> {
        let driver = self.live_driver().await?;

        let handle = driver
            .new_tab()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to open new tab: {}", e)))?;

        Ok(handle.to_string())
    }

    /// Open a new window.
    pub async fn new_window(&self) -> Result<String> {
        let driver = self.live_driver().await?;

        let handle = driver
            .new_window()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to open new window: {}", e)))?;

        Ok(handle.to_string())
    }

    /// Switch to a window by handle.
    pub async fn switch_to_window(&self, handle: &str) -> Result<()> {
        let driver = self.live_driver().await?;

        let window_handle = thirtyfour::WindowHandle::from(handle.to_string());
        driver
            .switch_to_window(window_handle)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to switch window: {}", e)))
    }

    /// Close the current window.
    pub async fn close_window(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .close_window()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to close window: {}", e)))
    }

    /// Maximize window.
    pub async fn maximize_window(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .maximize_window()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to maximize window: {}", e)))
    }

    /// Minimize window.
    pub async fn minimize_window(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .minimize_window()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to minimize window: {}", e)))
    }

    /// Set window size and position.
    pub async fn set_window_rect(&self, x: i32, y: i32, width: u32, height: u32) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .set_window_rect(x as i64, y as i64, width, height)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to set window rect: {}", e)))
    }

    // =========================================================================
    // Frame Management
    // =========================================================================

    /// Switch to frame by index.
    pub async fn switch_to_frame_by_index(&self, index: u16) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .enter_frame(index)
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to switch to frame {}: {}", index, e)))
    }

    /// Switch to frame by selector.
    pub async fn switch_to_frame_by_selector(&self, selector: &str) -> Result<()> {
        let element = self.find_element(selector).await?;

        element
            .enter_frame()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to switch to frame: {}", e)))
    }

    /// Switch to parent frame.
    pub async fn switch_to_parent_frame(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .enter_parent_frame()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to switch to parent frame: {}", e)))
    }

    /// Switch to default content.
    pub async fn switch_to_default_content(&self) -> Result<()> {
        let driver = self.live_driver().await?;

        driver
            .enter_default_frame()
            .await
            .map_err(|e| AdkError::Tool(format!("Failed to switch to default content: {}", e)))
    }

    // =========================================================================
    // Advanced Actions
    // =========================================================================

    /// Drag and drop.
    pub async fn drag_and_drop(&self, source_selector: &str, target_selector: &str) -> Result<()> {
        let source = self.find_element(source_selector).await?;
        let target = self.find_element(target_selector).await?;

        // Use JavaScript-based drag and drop for broader compatibility
        source
            .js_drag_to(&target)
            .await
            .map_err(|e| AdkError::Tool(format!("Drag and drop failed: {}", e)))
    }

    /// Right-click (context click).
    pub async fn right_click(&self, selector: &str) -> Result<()> {
        let escaped = escape_js_string(selector);
        let script = format!(
            r#"
            var element = document.querySelector('{escaped}');
            if (element) {{
                var event = new MouseEvent('contextmenu', {{
                    'view': window,
                    'bubbles': true,
                    'cancelable': true
                }});
                element.dispatchEvent(event);
                return true;
            }}
            return false;
            "#,
        );

        let result = self.execute_script(&script).await?;
        if result.as_bool() != Some(true) {
            return Err(AdkError::Tool(format!("Element not found: {}", selector)));
        }
        Ok(())
    }

    /// Focus an element.
    pub async fn focus_element(&self, selector: &str) -> Result<()> {
        let element = self.find_element(selector).await?;
        element.focus().await.map_err(|e| AdkError::Tool(format!("Focus failed: {}", e)))
    }

    /// Get element state.
    pub async fn get_element_state(&self, selector: &str) -> Result<ElementState> {
        let element = self.find_element(selector).await?;

        let is_displayed = element.is_displayed().await.unwrap_or(false);
        let is_enabled = element.is_enabled().await.unwrap_or(false);
        let is_selected = element.is_selected().await.unwrap_or(false);
        let is_clickable = element.is_clickable().await.unwrap_or(false);

        Ok(ElementState { is_displayed, is_enabled, is_selected, is_clickable })
    }

    /// Press a key, optionally with modifier keys (Ctrl, Alt, Shift, Meta).
    pub async fn press_key(
        &self,
        key: &str,
        selector: Option<&str>,
        modifiers: &[&str],
    ) -> Result<()> {
        // Map modifier names to WebDriver key codes
        let modifier_codes: Vec<&str> = modifiers
            .iter()
            .filter_map(|m| match m.to_lowercase().as_str() {
                "ctrl" | "control" => Some("\u{E009}"),
                "alt" => Some("\u{E00A}"),
                "shift" => Some("\u{E008}"),
                "meta" | "command" | "cmd" => Some("\u{E03D}"),
                _ => None,
            })
            .collect();

        let key_str = match key.to_lowercase().as_str() {
            "enter" => "\u{E007}",
            "tab" => "\u{E004}",
            "escape" | "esc" => "\u{E00C}",
            "backspace" => "\u{E003}",
            "delete" => "\u{E017}",
            "arrowup" | "up" => "\u{E013}",
            "arrowdown" | "down" => "\u{E015}",
            "arrowleft" | "left" => "\u{E012}",
            "arrowright" | "right" => "\u{E014}",
            "home" => "\u{E011}",
            "end" => "\u{E010}",
            "pageup" => "\u{E00E}",
            "pagedown" => "\u{E00F}",
            "space" => " ",
            _ => key,
        };

        // Build the full key sequence: modifiers down + key + modifiers up
        let mut key_sequence = String::new();
        for code in &modifier_codes {
            key_sequence.push_str(code);
        }
        key_sequence.push_str(key_str);
        // Release modifiers in reverse order
        for code in modifier_codes.iter().rev() {
            key_sequence.push_str(code);
        }

        if let Some(sel) = selector {
            let element = self.find_element(sel).await?;
            element
                .send_keys(&key_sequence)
                .await
                .map_err(|e| AdkError::Tool(format!("Key press failed: {e}")))?;
        } else {
            // Send to active element via WebDriver
            let driver = self.live_driver().await?;

            let active = driver
                .active_element()
                .await
                .map_err(|e| AdkError::Tool(format!("No active element: {e}")))?;
            active
                .send_keys(&key_sequence)
                .await
                .map_err(|e| AdkError::Tool(format!("Key press failed: {e}")))?;
        }

        Ok(())
    }

    /// Upload a file.
    pub async fn upload_file(&self, selector: &str, file_path: &str) -> Result<()> {
        let element = self.find_element(selector).await?;
        element
            .send_keys(file_path)
            .await
            .map_err(|e| AdkError::Tool(format!("File upload failed: {}", e)))
    }

    /// Print page to PDF.
    pub async fn print_to_pdf(&self, landscape: bool, scale: f64) -> Result<String> {
        let driver = self.live_driver().await?;

        let params = PrintParameters {
            orientation: if landscape {
                PrintOrientation::Landscape
            } else {
                PrintOrientation::Portrait
            },
            scale,
            ..Default::default()
        };

        driver
            .print_page_base64(params)
            .await
            .map_err(|e| AdkError::Tool(format!("Print to PDF failed: {}", e)))
    }

    /// Build browser capabilities based on configuration.
    fn build_capabilities(&self) -> Result<Capabilities> {
        let caps = match self.config.browser {
            BrowserType::Chrome => {
                let mut caps = DesiredCapabilities::chrome();
                if self.config.headless {
                    caps.add_arg("--headless=new").map_err(|e| {
                        AdkError::Tool(format!("Failed to add headless arg: {}", e))
                    })?;
                }
                caps.add_arg("--no-sandbox")
                    .map_err(|e| AdkError::Tool(format!("Failed to add no-sandbox: {}", e)))?;
                caps.add_arg("--disable-dev-shm-usage")
                    .map_err(|e| AdkError::Tool(format!("Failed to add disable-dev-shm: {}", e)))?;

                if let Some(ref ua) = self.config.user_agent {
                    caps.add_arg(&format!("--user-agent={}", ua))
                        .map_err(|e| AdkError::Tool(format!("Failed to add user-agent: {}", e)))?;
                }

                for arg in &self.config.browser_args {
                    caps.add_arg(arg).map_err(|e| {
                        AdkError::Tool(format!("Failed to add arg '{}': {}", arg, e))
                    })?;
                }

                caps.into()
            }
            BrowserType::Firefox => {
                let mut caps = DesiredCapabilities::firefox();
                if self.config.headless {
                    caps.add_arg("-headless")
                        .map_err(|e| AdkError::Tool(format!("Failed to add headless: {}", e)))?;
                }
                caps.into()
            }
            BrowserType::Safari => DesiredCapabilities::safari().into(),
            BrowserType::Edge => {
                let mut caps = DesiredCapabilities::edge();
                if self.config.headless {
                    caps.add_arg("--headless")
                        .map_err(|e| AdkError::Tool(format!("Failed to add headless: {}", e)))?;
                }
                caps.into()
            }
        };

        Ok(caps)
    }
}

impl Drop for BrowserSession {
    fn drop(&mut self) {
        // Note: Can't do async cleanup in Drop, but thirtyfour handles this gracefully
        tracing::debug!("BrowserSession dropped");
    }
}

/// Create a shared browser session.
pub fn shared_session(config: BrowserConfig) -> Arc<BrowserSession> {
    Arc::new(BrowserSession::new(config))
}

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

    #[test]
    fn test_session_creation() {
        let session = BrowserSession::with_defaults();
        assert!(!session.config().headless || session.config().headless); // Always true, just testing creation
    }

    #[tokio::test]
    async fn test_session_not_started() {
        // Use a bogus WebDriver URL so ensure_started's auto-start attempt fails
        let config = BrowserConfig::new().webdriver_url("http://127.0.0.1:1");
        let session = BrowserSession::new(config);
        assert!(!session.is_active().await);

        // navigate() calls live_driver() → ensure_started() → start(),
        // which fails because the WebDriver is unreachable.
        let result = session.navigate("https://example.com").await;
        assert!(result.is_err());
    }

    #[test]
    fn test_build_capabilities_chrome_headless() {
        let config = BrowserConfig::new().headless(true);
        let session = BrowserSession::new(config);
        let caps = session.build_capabilities();
        assert!(caps.is_ok());
    }

    #[test]
    fn test_build_capabilities_firefox() {
        let config = BrowserConfig::new().browser(BrowserType::Firefox);
        let session = BrowserSession::new(config);
        let caps = session.build_capabilities();
        assert!(caps.is_ok());
    }

    #[test]
    fn test_build_capabilities_safari() {
        let config = BrowserConfig::new().browser(BrowserType::Safari);
        let session = BrowserSession::new(config);
        let caps = session.build_capabilities();
        assert!(caps.is_ok());
    }

    #[test]
    fn test_build_capabilities_edge() {
        let config = BrowserConfig::new().browser(BrowserType::Edge);
        let session = BrowserSession::new(config);
        let caps = session.build_capabilities();
        assert!(caps.is_ok());
    }

    #[test]
    fn test_build_capabilities_with_user_agent() {
        let config = BrowserConfig::new().user_agent("CustomAgent/1.0");
        let session = BrowserSession::new(config);
        let caps = session.build_capabilities();
        assert!(caps.is_ok());
    }

    #[test]
    fn test_build_capabilities_with_extra_args() {
        let config = BrowserConfig::new().add_arg("--disable-gpu").add_arg("--window-size=800,600");
        let session = BrowserSession::new(config);
        let caps = session.build_capabilities();
        assert!(caps.is_ok());
    }
}