"""
Screenshot Tool

This tool provides screenshot capabilities for capturing web pages or desktop.
Note: This is a simplified implementation that would need browser automation
in a real environment.
"""

import os
import time
from typing import Dict, Any, List
from tools.tool_registry import BaseTool


class ScreenshotTool(BaseTool):
    """Tool for taking screenshots"""
    
    def __init__(self):
        super().__init__(
            name="screenshot_tool",
            description="Capture screenshots of web pages or desktop. Supports full page screenshots and element-specific captures."
        )
        self.screenshot_dir = "./screenshots"
        os.makedirs(self.screenshot_dir, exist_ok=True)
    
    def execute(self, query: str, context: Dict[str, Any]) -> Any:
        """
        Execute screenshot capture based on query
        
        Args:
            query: User query containing screenshot instructions
            context: Task analysis and other context
            
        Returns:
            Screenshot capture result
        """
        # Extract URL if present
        url = self._extract_url(query)
        
        if url:
            return self._capture_webpage_screenshot(url, query)
        else:
            return self._capture_desktop_screenshot(query)
    
    def _extract_url(self, query: str) -> str:
        """Extract URL from query"""
        import re
        
        # Look for URLs in the query
        url_pattern = r'https?://[^\s<>"{}|\\^`\[\]]+'
        urls = re.findall(url_pattern, query)
        
        if urls:
            return urls[0]
        
        # Look for domain names
        domain_pattern = r'\b(?:[a-zA-Z0-9-]+\.)+[a-zA-Z]{2,}\b'
        domains = re.findall(domain_pattern, query)
        
        if domains:
            domain = domains[0]
            if not domain.startswith(('http://', 'https://')):
                return f"https://{domain}"
            return domain
        
        return None
    
    def _capture_webpage_screenshot(self, url: str, query: str) -> Dict[str, Any]:
        """Capture screenshot of a webpage using Playwright if available, else install and use it, else fallback to instructions."""
        import subprocess
        import sys
        try:
            try:
                from playwright.sync_api import sync_playwright
            except ImportError:
                subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'playwright'])
                try:
                    subprocess.check_call([sys.executable, '-m', 'playwright', 'install', 'chromium'])
                except Exception:
                    subprocess.check_call(['playwright', 'install', 'chromium'])
                from playwright.sync_api import sync_playwright
            # Double-check browser binaries
            try:
                # Try launching browser to see if binaries are present
                with sync_playwright() as p:
                    browser = p.chromium.launch()
                    browser.close()
            except Exception as e:
                # Try to install again using CLI
                try:
                    subprocess.check_call(['playwright', 'install', 'chromium'])
                except Exception as cli_e:
                    return {
                        "success": False,
                        "url": url,
                        "error": f"Playwright install failed: {e}; CLI install also failed: {cli_e}",
                        "message": f"Could not install browser binaries for Playwright. Please install manually.",
                    }
            timestamp = int(time.time())
            filename = f"webpage_screenshot_{timestamp}.png"
            filepath = os.path.join(self.screenshot_dir, filename)
            # Actually capture the screenshot
            with sync_playwright() as p:
                browser = p.chromium.launch()
                page = browser.new_page()
                page.goto(url)
                page.screenshot(path=filepath)
                browser.close()
            return {
                "success": True,
                "url": url,
                "filepath": filepath,
                "message": f"Screenshot captured and saved to {filepath}"
            }
        except Exception as e:
            return {
                "success": False,
                "url": url,
                "error": str(e),
                "message": f"Error capturing screenshot of {url}: {str(e)}",
                "suggested_implementation": {
                    "tool": "playwright",
                    "code_example": f"""
from playwright.sync_api import sync_playwright

def capture_screenshot(url, filepath):
    with sync_playwright() as p:
        browser = p.chromium.launch()
        page = browser.new_page()
        page.goto('{url}')
        page.screenshot(path='{filepath}')
        browser.close()
    return filepath
"""
                }
            }
    
    def _capture_desktop_screenshot(self, query: str) -> Dict[str, Any]:
        """Capture desktop screenshot"""
        try:
            # This would require desktop screenshot capabilities
            # Using libraries like PIL (Pillow) with ImageGrab
            
            timestamp = int(time.time())
            filename = f"desktop_screenshot_{timestamp}.png"
            filepath = os.path.join(self.screenshot_dir, filename)
            
            # Simulate desktop screenshot
            result = {
                "success": False,
                "error": "Desktop screenshot not available in this environment",
                "message": "Desktop screenshot capture requires appropriate permissions and libraries",
                "suggested_implementation": {
                    "tool": "pillow",
                    "code_example": f"""
from PIL import ImageGrab

def capture_desktop_screenshot(filepath):
    screenshot = ImageGrab.grab()
    screenshot.save('{filepath}')
    return filepath
"""
                }
            }
            
            return result
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "message": f"Error capturing desktop screenshot: {str(e)}"
            }
    
    def get_capabilities(self) -> List[str]:
        """Return capabilities of this tool"""
        return [
            "webpage_screenshot",
            "desktop_screenshot", 
            "visual_capture",
            "image_generation"
        ]
    
    def validate_input(self, query: str, context: Dict[str, Any]) -> bool:
        """Validate input for screenshot operations"""
        screenshot_keywords = [
            'screenshot', 'capture', 'snap', 'image', 'picture',
            'screen', 'grab', 'save'
        ]
        
        query_lower = query.lower()
        return any(keyword in query_lower for keyword in screenshot_keywords)
    
    def list_screenshots(self) -> List[Dict[str, Any]]:
        """List all captured screenshots"""
        screenshots = []
        
        if os.path.exists(self.screenshot_dir):
            for filename in os.listdir(self.screenshot_dir):
                if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                    filepath = os.path.join(self.screenshot_dir, filename)
                    file_stat = os.stat(filepath)
                    
                    screenshots.append({
                        "filename": filename,
                        "filepath": filepath,
                        "size": file_stat.st_size,
                        "created": file_stat.st_ctime,
                        "modified": file_stat.st_mtime
                    })
        
        return screenshots
    
    def get_screenshot_directory(self) -> str:
        """Get the screenshot directory path"""
        return self.screenshot_dir

