ditto-os 0.1.0

use std::collections::HashMap;
use std::sync::Arc;
use std::path::Path;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use uuid::Uuid;
use chromiumoxide::{Browser, Page};
use chromiumoxide::cdp::browser_protocol::page;

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()
    }
}

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...");
        // This method can be used to pre-initialize chromiumoxide connections
        // For now, it's a no-op since we create browsers on-demand
        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
        );

        // Start real chromiumoxide instance
        let browser = match browser_type {
            BrowserType::Chrome => {
                let (browser, mut handler) = chromiumoxide::Browser::builder()
                    .with_headless(true)
                    .build()?;

                // Start the handler task
                tokio::spawn(async move {
                    loop {
                        let _ = handler.recv().await;
                    }
                });

                browser
            }
            _ => return Err(anyhow::anyhow!("Unsupported browser type: {:?}", browser_type)),
        };

        // Create a new page
        let page = browser.new_page("about:blank").await?;

        // Store browser and page instances
        {
            let mut browsers = self.browsers.write();
            browsers.insert(session_id.clone(), browser);
            let mut pages = self.pages.write();
            pages.insert(session_id.clone(), page);
        }

        let session = BrowserSession {
            id: session_id.clone(),
            browser_type,
            agent_id,
            status: SessionStatus::Starting,
            created_at: chrono::Utc::now(),
            last_activity: chrono::Utc::now(),
            url: None,
            title: None,
            chromiumoxide_handle: Some(format!("chromiumoxide_{}", session_id)),
        };

        // Add session to manager
        {
            let mut sessions = self.sessions.write();
            sessions.insert(session_id.clone(), session);
        }

        // Navigate to about:blank to start
        if let Err(e) = self.start_browser_container(&session_id).await {
            error!("Failed to start browser container for {}: {}", session_id, e);
            // Clean up on failure
            {
                let mut browsers = self.browsers.write();
                browsers.remove(&session_id);
                let mut pages = self.pages.write();
                pages.remove(&session_id);
            }
            let mut sessions = self.sessions.write();
            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().clone();
        sessions.get(session_id).cloned()
    }

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

        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();

        // Check if session exists and is running
        let sessions = self.sessions.read().clone();
        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));
        }

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

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

        let result = match &command.action {
            BrowserAction::Navigate { url } => {
                // Real navigation using chromiumoxide
                page.goto(url).await?;
                
                // Get actual page title
                let title = page.get_title().await?;
                
                // Update session info
                {
                    let mut sessions = self.sessions.write();
                    if let Some(session) = sessions.get_mut(session_id) {
                        session.url = Some(url.clone());
                        session.title = Some(title.clone());
                    }
                }

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

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

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

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

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

                Ok(serde_json::json!({
                    "title": title
                }))
            }
            BrowserAction::GetUrl {} => {
                // Get real current URL
                let url = page.get_url().await?;

                Ok(serde_json::json!({
                    "url": url
                }))
            }
            BrowserAction::Refresh {} => {
                // Real page refresh
                page.reload().await?;

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

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

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

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

                Ok(serde_json::json!({
                    "action": "get_attribute",
                    "selector": selector,
                    "attribute": attribute,
                    "value": attr_value
                }))
            }
            BrowserAction::Screenshot { path } => {
                // Take screenshot using chromiumoxide
                let screenshot_path = path.unwrap_or_else(|| format!("screenshot_{}.png", session_id));
                page.save_screenshot(Path::new(&screenshot_path), CaptureScreenshotFormat::Png).await?;

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

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

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

        // Update last activity
        {
            let mut sessions = self.sessions.write();
            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);

        // Close the chromiumoxide browser connection
        {
            let mut browsers = self.browsers.write().await;
            let mut pages = self.pages.write().await;
            if let Some(_browser) = browsers.remove(session_id) {
                pages.remove(session_id);
                info!("Chromiumoxide browser closed for session: {}", session_id);
            }
        }

        // Remove session from memory
        {
            let mut sessions = self.sessions.write();
            sessions.remove(session_id);
        }

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

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

        let pages = self.pages.read().await.clone();
        if let Some(page) = pages.get(session_id) {
            // Navigate to about:blank to start
            page.goto("about:blank").await?;
            
            // Update session status to running
            {
                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);
        } else {
            return Err(anyhow::anyhow!("Page not found for session: {}", session_id));
        }

        Ok(())
    }

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

        // Close the chromiumoxide browser connection
        {
            let mut browsers = self.browsers.write().await;
            let mut pages = self.pages.write().await;
            if let Some(_browser) = 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: std::collections::HashMap<String, BrowserSession> = 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().clone();
            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: std::collections::HashMap<String, BrowserSession> = self.sessions.read().await.clone();
        
        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() {
            // Count browser types
            let browser_type_str = format!("{:?}", session.browser_type);
            *browser_type_counts.entry(browser_type_str).or_insert(0) += 1;
            
            // Count statuses
            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,
        })
    }
}