#!/usr/bin/env python3
# Hermes Node Protocol
# Copyright (c) 2026 Stefy (nextime) Lanza <stefy@nexlab.net>
# All rights reserved.
# 
# This software is released under the MIT License with a copyleft clause.
# See the LICENSE file for full terms.
"""
Browser control module for Hermes Node Protocol.
Implements the interface expected by hermes_node_agent.py.
"""

import asyncio
import base64
import logging
from typing import Dict, Any, Optional
from pathlib import Path

try:
    from playwright.async_api import async_playwright, Browser, BrowserContext, Page, Error as PlaywrightError
    HAS_PLAYWRIGHT = True
except ImportError:
    HAS_PLAYWRIGHT = False
    Browser = BrowserContext = Page = None

logger = logging.getLogger(__name__)

class BrowserController:
    """Manages browser instances and contexts for remote control."""
    
    def __init__(self):
        self.playwright = None
        self.browser: Optional[Browser] = None
        self.contexts: Dict[str, BrowserContext] = {}  # name -> context
        self.default_context: Optional[BrowserContext] = None
        self.pages: Dict[str, Page] = {}  # page_id -> page
        self.lock = asyncio.Lock()
        
    async def initialize(self):
        """Initialize Playwright."""
        if not HAS_PLAYWRIGHT:
            raise RuntimeError("Playwright not installed. Run: pip install playwright && playwright install chromium")
        
        self.playwright = await async_playwright().start()
        logger.info("Playwright initialized")
    
    async def shutdown(self):
        """Clean shutdown of all browser resources."""
        try:
            # Close all contexts
            for ctx in self.contexts.values():
                await ctx.close()
            if self.default_context:
                await self.default_context.close()
            
            # Close browser
            if self.browser:
                await self.browser.close()
            
            # Stop playwright
            if self.playwright:
                await self.playwright.stop()
            
            logger.info("Browser controller shutdown complete")
        except Exception as e:
            logger.error(f"Error during shutdown: {e}")
    
    # Launch and context management
    
    async def launch(self, config: Dict = None) -> Dict[str, Any]:
        """Launch the browser."""
        config = config or {}
        headless = config.get("headless", False)
        attach = config.get("attach", False)
        cdp_url = config.get("cdp_url", "http://localhost:9222")
        browser_type = config.get("browser_type", "chromium")
        
        try:
            if attach:
                # Connect to existing browser
                endpoints = {
                    "chromium": f"{cdp_url}",
                    "firefox": f"{cdp_url.replace(':9222', ':9200') if ':9222' in cdp_url else cdp_url}",
                    "webkit": f"{cdp_url.replace(':9222', ':9223') if ':9222' in cdp_url else cdp_url}"
                }
                endpoint = endpoints.get(browser_type, cdp_url)
                
                if browser_type == "chromium":
                    self.browser = await self.playwright.chromium.connect_over_cdp(endpoint)
                elif browser_type == "firefox":
                    self.browser = await self.playwright.firefox.connect_over_cdp(endpoint)
                elif browser_type == "webkit":
                    self.browser = await self.playwright.webkit.connect_over_cdp(endpoint)
                else:
                    self.browser = await self.playwright.chromium.connect_over_cdp(endpoint)
                
                logger.info(f"Attached to existing {browser_type} browser at {endpoint}")
            else:
                # Launch new browser
                browser_types = {
                    "chromium": self.playwright.chromium,
                    "firefox": self.playwright.firefox,
                    "webkit": self.playwright.webkit
                }
                launch_browser = browser_types.get(browser_type, self.playwright.chromium)
                
                args = config.get("args", ['--no-sandbox', '--disable-setuid-sandbox'])
                
                # Add extension paths if provided (Chromium only)
                extension_paths = config.get("extension_paths", [])
                if extension_paths and browser_type == "chromium":
                    args.append('--load-extension=' + ','.join(extension_paths))
                
                self.browser = await launch_browser.launch(
                    headless=headless,
                    args=args
                )
                logger.info(f"Launched {browser_type} browser (headless={headless})")
            
            return {"success": True, "browser": browser_type, "headless": headless, "mode": "attach" if attach else "launch"}
            
        except Exception as e:
            logger.error(f"Failed to launch/attach browser: {e}")
            return {"success": False, "error": str(e)}
    
    async def create_context(self, config: Dict = None) -> Dict[str, Any]:
        """Create a new browser context."""
        if not self.browser:
            return {"success": False, "error": "Browser not launched"}
        
        config = config or {}
        context_name = config.get("name")
        
        async with self.lock:
            try:
                # Check for incognito
                incognito = config.get("incognito", False)
                if incognito:
                    ctx = await self.browser.new_context()
                    ctx_name = f"_incognito_{id(ctx)}"
                else:
                    ctx = await self.browser.new_context()
                    ctx_name = context_name or f"ctx_{len(self.contexts)}"
                
                self.contexts[ctx_name] = ctx
                
                # Create initial page if requested
                create_page = config.get("create_page", True)
                page_id = None
                if create_page:
                    page = await ctx.new_page()
                    page_id = f"page_{len(self.pages)}"
                    self.pages[page_id] = page
                
                logger.info(f"Created context: {ctx_name} (incognito={incognito})")
                return {
                    "success": True,
                    "context_name": ctx_name,
                    "page_id": page_id,
                    "incognito": incognito
                }
            except Exception as e:
                logger.error(f"Failed to create context: {e}")
                return {"success": False, "error": str(e)}
    
    async def new_page(self, context_name: Optional[str] = None) -> Dict[str, Any]:
        """Create a new page in specified context (or default)."""
        try:
            ctx = await self._get_context(context_name)
            page = await ctx.new_page()
            page_id = f"page_{len(self.pages)}"
            self.pages[page_id] = page
            
            return {"success": True, "page_id": page_id, "context": context_name or "default"}
        except Exception as e:
            logger.error(f"Failed to create page: {e}")
            return {"success": False, "error": str(e)}
    
    # Navigation
    
    async def navigate(self, page_id: str, url: str, wait_until: str = "load", 
                       timeout: int = 30000) -> Dict[str, Any]:
        """Navigate to a URL."""
        page = await self._get_page(page_id)
        await page.goto(url, wait_until=wait_until, timeout=timeout)
        
        return {
            "success": True,
            "url": page.url,
            "title": await page.title()
        }
    
    # Interaction
    
    async def click(self, page_id: str, selector: str, timeout: int = 30000) -> Dict[str, Any]:
        """Click an element."""
        page = await self._get_page(page_id)
        await page.click(selector, timeout=timeout)
        return {"success": True, "selector": selector}
    
    async def fill(self, page_id: str, selector: str, value: str, timeout: int = 30000) -> Dict[str, Any]:
        """Fill a form field."""
        page = await self._get_page(page_id)
        await page.fill(selector, value, timeout=timeout)
        return {"success": True, "selector": selector, "value": value}
    
    async def type_text(self, page_id: str, selector: str, text: str, 
                        delay: int = 0, timeout: int = 30000) -> Dict[str, Any]:
        """Type text character by character."""
        page = await self._get_page(page_id)
        await page.type(selector, text, delay=delay, timeout=timeout)
        return {"success": True, "selector": selector, "text": text}
    
    async def wait_for_selector(self, page_id: str, selector: str, 
                               state: str = "visible", timeout: int = 30000) -> Dict[str, Any]:
        """Wait for an element to be in a state."""
        page = await self._get_page(page_id)
        await page.wait_for_selector(selector, state=state, timeout=timeout)
        return {"success": True, "selector": selector, "state": state}
    
    # Script execution
    
    async def execute_script(self, page_id: str, script: str) -> Dict[str, Any]:
        """Execute JavaScript in page context (no return)."""
        page = await self._get_page(page_id)
        await page.evaluate(script)
        return {"success": True}
    
    async def evaluate(self, page_id: str, expression: str, arg: Any = None) -> Dict[str, Any]:
        """Evaluate JavaScript and return result."""
        page = await self._get_page(page_id)
        result = await page.evaluate(expression, arg)
        return {"success": True, "result": result}
    
    # Inspection
    
    async def screenshot(self, page_id: str, full_page: bool = False, 
                        path: Optional[str] = None) -> Dict[str, Any]:
        """Take a screenshot."""
        page = await self._get_page(page_id)
        
        if path:
            await page.screenshot(path=path, full_page=full_page)
            return {"success": True, "path": path, "format": "png"}
        else:
            screenshot_bytes = await page.screenshot(full_page=full_page)
            screenshot_b64 = base64.b64encode(screenshot_bytes).decode('utf-8')
            return {"success": True, "screenshot": screenshot_b64, "format": "png", "encoding": "base64"}
    
    async def get_content(self, page_id: str) -> Dict[str, Any]:
        """Get page HTML content."""
        page = await self._get_page(page_id)
        content = await page.content()
        return {"success": True, "content": content}
    
    async def get_title(self, page_id: str) -> Dict[str, Any]:
        """Get page title."""
        page = await self._get_page(page_id)
        title = await page.title()
        return {"success": True, "title": title}
    
    async def get_url(self, page_id: str) -> Dict[str, Any]:
        """Get current page URL."""
        page = await self._get_page(page_id)
        return {"success": True, "url": page.url}
    
    # Cleanup
    
    async def close_page(self, page_id: str) -> Dict[str, Any]:
        """Close a specific page."""
        if page_id not in self.pages:
            return {"success": False, "error": f"Page not found: {page_id}"}
        
        page = self.pages[page_id]
        await page.close()
        del self.pages[page_id]
        
        return {"success": True, "page_id": page_id}
    
    async def close_context(self, context_name: str) -> Dict[str, Any]:
        """Close a named context and all its pages."""
        if context_name not in self.contexts:
            return {"success": False, "error": f"Context not found: {context_name}"}
        
        ctx = self.contexts[context_name]
        await ctx.close()
        del self.contexts[context_name]
        
        # Clean up any pages from this context
        self.pages = {pid: p for pid, p in self.pages.items() 
                      if p.context.name != context_name}
        
        logger.info(f"Closed context: {context_name}")
        return {"success": True, "context": context_name}
    # Browser extension support
    
    async def load_extension(self, extension_path: str) -> Dict[str, Any]:
        """Load a Chrome/Chromium unpacked extension."""
        if not self.browser:
            return {"success": False, "error": "Browser not launched"}
        
        if not Path(extension_path).exists():
            return {"success": False, "error": f"Extension path not found: {extension_path}"}
        
        # Extension is loaded via args during browser launch, this reloads it
        return {"success": False, "error": "Extension must be passed in launch config with 'extension_paths'"}
    
    async def execute_extension_script(self, extension_id: str, script: str) -> Dict[str, Any]:
        """Execute script in extension context via CDP."""
        # This would use CDP to inject into extension background page
        return {"success": False, "error": "Use CDP layer directly with Runtime.evaluate"}
    
    async def list_extensions(self) -> Dict[str, Any]:
        """List installed extensions via CDP."""
        return {"success": False, "error": "Use CDP: Browser.getBrowserCommandLine and querying"}

    async def close(self) -> Dict[str, Any]:
        """Close all contexts and the browser."""
        await self.shutdown()
        self.contexts = {}
        self.pages = {}
        self.default_context = None
        self.browser = None
        return {"success": True, "message": "Browser closed"}
    
    # Listing
    
    async def list_pages(self) -> Dict[str, Any]:
        """List all active pages."""
        pages_info = []
        for page_id, page in self.pages.items():
            pages_info.append({
                "page_id": page_id,
                "url": page.url,
                "title": await page.title()
            })
        return {"success": True, "pages": pages_info}
    
    async def list_contexts(self) -> Dict[str, Any]:
        """List all contexts."""
        contexts_info = []
        for name, ctx in self.contexts.items():
            contexts_info.append({
                "name": name,
                "pages_count": len(ctx.pages)
            })
        # Include default context if exists
        if self.default_context:
            contexts_info.append({
                "name": "default",
                "pages_count": len(self.default_context.pages)
            })
        return {"success": True, "contexts": contexts_info}
    
    # Advanced APIs
    
    async def playwright_command(self, page_id: str, command: str, 
                                args: list = None, kwargs: dict = None) -> Dict[str, Any]:
        """Execute arbitrary Playwright API."""
        args = args or []
        kwargs = kwargs or {}
        page = await self._get_page(page_id)
        
        # Parse command path: "locator.click" -> page.locator(...).click(...)
        parts = command.split(".")
        obj = page
        
        for i, part in enumerate(parts[:-1]):
            # Get the attribute
            attr = getattr(obj, part)
            
            # If it's a method, we need to call it with an arg from args
            # Pattern: part.method(args[i])
            if callable(attr):
                if args:
                    obj = attr(args[0])
                    args = args[1:]
                else:
                    obj = attr()
            else:
                obj = attr
        
        # Call final method
        final_method = getattr(obj, parts[-1])
        if callable(final_method):
            result = await final_method(*args, **kwargs)
        else:
            result = final_method
        
        return {"success": True, "result": result}
    
    async def cdp_command(self, page_id: str, method: str, params: dict = None) -> Dict[str, Any]:
        """Send raw CDP command."""
        page = await self._get_page(page_id)
        cdp = await page.context.new_cdp_session(page)
        result = await cdp.send(method, params or {})
        return {"success": True, "result": result}
    
    # Helpers
    
    async def _get_context(self, context_name: Optional[str] = None) -> BrowserContext:
        """Get context by name or default."""
        if context_name:
            if context_name not in self.contexts:
                raise ValueError(f"Context not found: {context_name}")
            return self.contexts[context_name]
        
        if not self.default_context:
            self.default_context = await self.browser.new_context()
        return self.default_context
    
    async def _get_page(self, page_id: str) -> Page:
        """Get page by ID."""
        if page_id not in self.pages:
            raise ValueError(f"Page not found: {page_id}")
        return self.pages[page_id]
