PicoBot/src/tools/browser/manager.rs

use std::collections::HashMap;
use std::path::{Component, Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{Result, anyhow, bail};
use tokio::sync::Mutex;
use uuid::Uuid;

use super::action::BrowserAction;
use super::runner::AgentBrowserRunner;
use super::security::validate_navigation;
use crate::bus::MediaRef;
use crate::config::{BrowserConfig, expand_path};
use crate::tools::{ToolResult, ToolResultWithMedia};

struct BrowserSession {
    agent_browser_id: String,
    gate: Mutex<()>,
    last_used: std::sync::Mutex<Instant>,
}

pub(super) struct BrowserManager {
    runner: AgentBrowserRunner,
    sessions: Mutex<HashMap<String, Arc<BrowserSession>>>,
    max_sessions: usize,
    idle_timeout: Duration,
    artifact_dir: PathBuf,
    allow_private_hosts: bool,
    allowed_domains: Vec<String>,
}

impl BrowserManager {
    pub(super) fn new(config: &BrowserConfig, workspace_dir: PathBuf) -> Result<Self> {
        if config.max_sessions == 0 {
            bail!("browser.max_sessions must be greater than zero");
        }
        if config.command.trim().is_empty() {
            bail!("browser.command cannot be empty");
        }
        let artifact_dir = expand_path(&config.artifact_dir);
        let artifact_dir = if artifact_dir.is_absolute() {
            artifact_dir
        } else {
            workspace_dir.join(artifact_dir)
        };
        Ok(Self {
            runner: AgentBrowserRunner::new(config, workspace_dir),
            sessions: Mutex::new(HashMap::new()),
            max_sessions: config.max_sessions,
            idle_timeout: Duration::from_secs(config.idle_timeout_secs.max(1)),
            artifact_dir,
            allow_private_hosts: config.allow_private_hosts,
            allowed_domains: config.allowed_domains.clone(),
        })
    }

    pub(super) async fn execute(
        &self,
        picobot_session_id: &str,
        action: BrowserAction,
    ) -> Result<ToolResultWithMedia> {
        if let BrowserAction::Open { url } = &action {
            validate_navigation(url, self.allow_private_hosts, &self.allowed_domains)
                .await
                .map_err(anyhow::Error::msg)?;
        }

        if action.is_close() {
            return self.close(picobot_session_id).await;
        }

        let screenshot_path = match action.screenshot_filename() {
            Some(filename) => Some(self.prepare_screenshot_path(filename).await?),
            None => None,
        };
        let (session, stale) = self.session_for(picobot_session_id).await?;
        for stale_session in stale {
            let _ = self
                .runner
                .run(&stale_session, &["close".to_string()])
                .await;
        }

        let _gate = session.gate.lock().await;
        let path_string = screenshot_path
            .as_ref()
            .map(|path| path.to_string_lossy().into_owned());
        let commands = action.commands(path_string.as_deref());
        let mut last_response = None;
        for command in commands {
            last_response = Some(self.runner.run(&session.agent_browser_id, &command).await?);
        }
        *session
            .last_used
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = Instant::now();

        let response =
            last_response.ok_or_else(|| anyhow!("browser action produced no command"))?;
        let mut output = self.runner.render_response(&response);
        let mut media_refs = Vec::new();
        if let Some(path) = screenshot_path {
            let metadata = tokio::fs::metadata(&path)
                .await
                .map_err(|error| anyhow!("agent-browser did not create screenshot: {error}"))?;
            if !metadata.is_file() || metadata.len() == 0 {
                bail!("agent-browser created an empty screenshot");
            }
            let canonical = tokio::fs::canonicalize(&path).await.unwrap_or(path);
            let canonical = canonical.to_string_lossy().into_owned();
            output = format!("Screenshot saved: {canonical}\n{output}");
            media_refs.push(MediaRef {
                path: canonical,
                media_type: "image".to_string(),
            });
        }
        Ok(ToolResultWithMedia {
            result: ToolResult {
                success: true,
                output,
                error: None,
            },
            media_refs,
        })
    }

    async fn session_for(
        &self,
        picobot_session_id: &str,
    ) -> Result<(Arc<BrowserSession>, Vec<String>)> {
        let now = Instant::now();
        let mut sessions = self.sessions.lock().await;
        if let Some(session) = sessions.get(picobot_session_id) {
            *session
                .last_used
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner()) = now;
            return Ok((session.clone(), Vec::new()));
        }

        let mut stale_ids = Vec::new();
        sessions.retain(|_, session| {
            let last_used = *session
                .last_used
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            let idle = now.duration_since(last_used) >= self.idle_timeout;
            let keep = !idle || session.gate.try_lock().is_err();
            if !keep {
                stale_ids.push(session.agent_browser_id.clone());
            }
            keep
        });
        if sessions.len() >= self.max_sessions {
            bail!(
                "browser session limit reached ({}); close another dialog browser or wait for idle cleanup",
                self.max_sessions
            );
        }
        let session = Arc::new(BrowserSession {
            agent_browser_id: format!("picobot-{}", Uuid::new_v4().simple()),
            gate: Mutex::new(()),
            last_used: std::sync::Mutex::new(now),
        });
        sessions.insert(picobot_session_id.to_string(), session.clone());
        Ok((session, stale_ids))
    }

    async fn close(&self, picobot_session_id: &str) -> Result<ToolResultWithMedia> {
        let session = self.sessions.lock().await.remove(picobot_session_id);
        let Some(session) = session else {
            return Ok(ToolResult {
                success: true,
                output: "Browser session is already closed.".to_string(),
                error: None,
            }
            .into());
        };
        let _gate = session.gate.lock().await;
        let response = self
            .runner
            .run(&session.agent_browser_id, &["close".to_string()])
            .await?;
        Ok(ToolResult {
            success: true,
            output: self.runner.render_response(&response),
            error: None,
        }
        .into())
    }

    async fn prepare_screenshot_path(&self, requested: Option<&str>) -> Result<PathBuf> {
        tokio::fs::create_dir_all(&self.artifact_dir).await?;
        let filename = match requested {
            Some(requested) => {
                let path = Path::new(requested);
                if path.is_absolute()
                    || path
                        .components()
                        .any(|component| !matches!(component, Component::Normal(_)))
                {
                    bail!("screenshot path must be a filename without directory components");
                }
                let filename = path
                    .file_name()
                    .and_then(|name| name.to_str())
                    .ok_or_else(|| anyhow!("invalid screenshot filename"))?;
                if !filename.to_ascii_lowercase().ends_with(".png") {
                    bail!("screenshot filename must end in .png");
                }
                filename.to_string()
            }
            None => format!(
                "picobot-screenshot-{}-{}.png",
                chrono::Utc::now().format("%Y%m%dT%H%M%S"),
                &Uuid::new_v4().simple().to_string()[..8]
            ),
        };
        Ok(self.artifact_dir.join(filename))
    }
}