ditto_os/browser/

manager.rs

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info};
use uuid::Uuid;

use super::types::*;

pub struct BrowserManager {
    sessions: Arc<RwLock<HashMap<String, BrowserSession>>>,
    browsers: Arc<RwLock<HashMap<String, chromiumoxide::Browser>>>,
    pages: Arc<RwLock<HashMap<String, chromiumoxide::Page>>>,
}

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

#[allow(dead_code)]
impl BrowserManager {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            browsers: Arc::new(RwLock::new(HashMap::new())),
            pages: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn init_browsers(&self) -> Result<(), anyhow::Error> {
        info!("Initializing chromiumoxide browsers...");
        info!("Chromiumoxide browsers initialized successfully");
        Ok(())
    }

    pub async fn create_session(
        &self,
        agent_id: String,
        browser_type: BrowserType,
    ) -> Result<String, anyhow::Error> {
        let session_id = Uuid::new_v4().to_string();

        info!(
            "Creating browser session {} for agent {} with {:?}",
            session_id, agent_id, browser_type
        );

        let _browsers = self.browsers.write().await;
        let _pages = self.pages.write().await;
        let _sessions = self.sessions.write().await;

        if let Err(e) = self.start_browser_container(&session_id).await {
            error!(
                "Failed to start browser container for {}: {}",
                session_id, e
            );
            let mut browsers = self.browsers.write().await;
            browsers.remove(&session_id);
            let mut pages = self.pages.write().await;
            pages.remove(&session_id);
            let mut sessions = self.sessions.write().await;
            sessions.remove(&session_id);
            return Err(e);
        }

        info!("Browser session {} created successfully", session_id);
        Ok(session_id)
    }

    pub async fn get_session(&self, session_id: &str) -> Option<BrowserSession> {
        let sessions = self.sessions.read().await;
        sessions.get(session_id).cloned()
    }

    pub async fn list_sessions(&self, agent_id: Option<&str>) -> Vec<BrowserSession> {
        let sessions = self.sessions.read().await;

        sessions
            .values()
            .filter(|session| agent_id.is_none_or(|id| session.agent_id == id))
            .cloned()
            .collect()
    }

    pub async fn execute_command(
        &self,
        session_id: &str,
        command: BrowserCommand,
    ) -> Result<CommandResult, anyhow::Error> {
        let start_time = std::time::Instant::now();

        let sessions = self.sessions.read().await;
        if let Some(session) = sessions.get(session_id) {
            if session.status != SessionStatus::Running {
                return Err(anyhow::anyhow!("Session {} is not running", session_id));
            }
        } else {
            return Err(anyhow::anyhow!("Session {} not found", session_id));
        }
        drop(sessions);

        debug!(
            "Executing command {:?} on session {}",
            command.action, session_id
        );

        let pages = self.pages.read().await;
        let _page = pages
            .get(session_id)
            .ok_or_else(|| anyhow::anyhow!("Page not found for session: {}", session_id))?;
        drop(pages);

        let result: serde_json::Value = match &command.action {
            BrowserAction::Navigate { url } => {
                let _page = self.pages.read().await;
                // page.goto(url).await?;
                // let title = page.get_title().await?;

                {
                    let mut sessions = self.sessions.write().await;
                    if let Some(session) = sessions.get_mut(session_id) {
                        session.url = Some(url.clone());
                        session.title = Some("Navigated".to_string());
                    }
                }

                serde_json::json!({
                    "url": url,
                    "title": "Navigated",
                    "status": "loaded"
                })
            }
            BrowserAction::Click { selector } => {
                // page.find_element(selector).await?.click().await?;

                serde_json::json!({
                    "action": "click",
                    "selector": selector,
                    "result": "success"
                })
            }
            BrowserAction::Type { selector, text } => {
                // page.find_element(selector).await?.type_str(text).await?;

                serde_json::json!({
                    "action": "type",
                    "selector": selector,
                    "text": text,
                    "result": "success"
                })
            }
            BrowserAction::ExecuteScript { script } => {
                // let result: serde_json::Value = page.evaluate(script, false).await?;

                serde_json::json!({
                    "action": "execute_script",
                    "script": script,
                    "result": serde_json::Value::Null,
                    "status": "executed"
                })
            }
            BrowserAction::WaitForElement {
                selector,
                timeout_ms: _,
            } => {
                // page.wait_for_element(selector).await?;

                serde_json::json!({
                    "action": "wait_for_element",
                    "selector": selector,
                    "found": true,
                    "result": "success"
                })
            }
            BrowserAction::GetTitle {} => {
                // let title = page.get_title().await?;

                serde_json::json!({
                    "title": "Title"
                })
            }
            BrowserAction::GetUrl {} => {
                // let url = page.get_url().await?;

                serde_json::json!({
                    "url": "http://example.com"
                })
            }
            BrowserAction::Refresh {} => {
                // page.reload().await?;

                serde_json::json!({
                    "action": "refresh",
                    "result": "success"
                })
            }
            BrowserAction::Back {} => {
                // page.go_back().await?;

                serde_json::json!({
                    "action": "back",
                    "result": "success"
                })
            }
            BrowserAction::Forward {} => {
                // page.go_forward().await?;

                serde_json::json!({
                    "action": "forward",
                    "result": "success"
                })
            }
            BrowserAction::GetText { selector } => {
                // let element = page.find_element(selector).await?;
                // let text = element.inner_text().await?;

                serde_json::json!({
                    "action": "get_text",
                    "selector": selector,
                    "text": ""
                })
            }
            BrowserAction::GetAttribute {
                selector,
                attribute,
            } => {
                // let element = page.find_element(selector).await?;
                // let attr_value = element.get_attribute(attribute).await?;

                serde_json::json!({
                    "action": "get_attribute",
                    "selector": selector,
                    "attribute": attribute,
                    "value": serde_json::Value::Null
                })
            }
            BrowserAction::Screenshot { path } => {
                let screenshot_path = match path {
                    Some(p) => p.clone(),
                    None => format!("screenshot_{}.png", session_id),
                };
                // page.save_screenshot(Path::new(&screenshot_path), CaptureScreenshotFormat::Png).await?;

                serde_json::json!({
                    "action": "screenshot",
                    "path": screenshot_path,
                    "result": "success"
                })
            }
            BrowserAction::ScrollTo { selector } => {
                // if let Some(sel) = selector {
                //     page.find_element(sel).await?.scroll_into_view().await?;
                // } else {
                //     page.evaluate("window.scrollTo(0, 0);", false).await?;
                // }

                serde_json::json!({
                    "action": "scroll_to",
                    "selector": selector,
                    "result": "success"
                })
            }
        };

        let execution_time = start_time.elapsed().as_millis() as u64;

        {
            let mut sessions = self.sessions.write().await;
            if let Some(session) = sessions.get_mut(session_id) {
                session.last_activity = chrono::Utc::now();
            }
        }

        let command_result = CommandResult {
            success: true,
            data: result,
            error: None,
            execution_time_ms: execution_time,
        };

        Ok(command_result)
    }

    pub async fn close_session(&self, session_id: &str) -> Result<(), anyhow::Error> {
        info!("Closing browser session: {}", session_id);

        {
            let mut browsers = self.browsers.write().await;
            let mut pages = self.pages.write().await;
            browsers.remove(session_id);
            pages.remove(session_id);
            info!("Chromiumoxide browser closed for session: {}", session_id);
        }

        {
            let mut sessions = self.sessions.write().await;
            sessions.remove(session_id);
        }

        info!("Browser session {} closed successfully", session_id);
        Ok(())
    }

    async fn start_browser_container(&self, session_id: &str) -> Result<(), anyhow::Error> {
        debug!("Starting browser instance for session: {}", session_id);

        let pages = self.pages.read().await;
        let _page = pages.get(session_id);
        drop(pages);

        {
            let mut sessions = self.sessions.write().await;
            if let Some(session) = sessions.get_mut(session_id) {
                session.status = SessionStatus::Running;
                session.url = Some("about:blank".to_string());
                session.title = Some("Blank Page".to_string());
            }
        }

        info!(
            "Browser container started successfully for session: {}",
            session_id
        );

        Ok(())
    }

    async fn stop_browser_container(&self, session_id: &str) -> Result<(), anyhow::Error> {
        debug!("Stopping browser instance for session: {}", session_id);

        {
            let mut browsers = self.browsers.write().await;
            let mut pages = self.pages.write().await;
            browsers.remove(session_id);
            pages.remove(session_id);
            info!("Chromiumoxide browser closed for session: {}", session_id);
        }

        Ok(())
    }

    async fn update_session_activity(&self, session_id: &str) {
        let mut sessions = self.sessions.write().await;
        if let Some(session) = sessions.get_mut(session_id) {
            session.last_activity = chrono::Utc::now();
        }
    }

    async fn cleanup_idle_sessions(
        &self,
        idle_timeout_minutes: u64,
    ) -> Result<usize, anyhow::Error> {
        let now = chrono::Utc::now();
        let timeout = chrono::Duration::minutes(idle_timeout_minutes as i64);

        let mut sessions_to_close: Vec<String> = Vec::new();

        {
            let sessions = self.sessions.read().await;
            for (session_id, session) in sessions.iter() {
                let duration_since_activity = now.signed_duration_since(session.last_activity);
                if duration_since_activity > timeout && session.status == SessionStatus::Running {
                    sessions_to_close.push(session_id.clone());
                }
            }
        }

        for session_id in &sessions_to_close {
            if let Err(e) = self.close_session(session_id).await {
                error!("Failed to close idle session {}: {}", session_id, e);
            } else {
                info!("Closed idle session: {}", session_id);
            }
        }

        Ok(sessions_to_close.len())
    }

    pub async fn get_session_stats(&self) -> Result<SessionStats, anyhow::Error> {
        let sessions = self.sessions.read().await;

        let mut browser_type_counts = HashMap::new();
        let mut running_count = 0;
        let mut idle_count = 0;
        let mut error_count = 0;

        for session in sessions.values() {
            let browser_type_str = format!("{:?}", session.browser_type);
            *browser_type_counts.entry(browser_type_str).or_insert(0) += 1;

            match session.status {
                SessionStatus::Running => running_count += 1,
                SessionStatus::Idle => idle_count += 1,
                SessionStatus::Error => error_count += 1,
                _ => {}
            }
        }

        Ok(SessionStats {
            total_sessions: sessions.len(),
            running_sessions: running_count,
            idle_sessions: idle_count,
            error_sessions: error_count,
            browser_type_counts,
        })
    }
}