"""
Screenshot capture module for SneakyScanner.

Uses Playwright with Chromium to capture screenshots of discovered web services.
"""

import os
import logging
from pathlib import Path
from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeout


class ScreenshotCapture:
    """
    Handles webpage screenshot capture for web services discovered during scanning.

    Uses Playwright with Chromium in headless mode to capture viewport screenshots
    of HTTP and HTTPS services. Handles SSL certificate errors gracefully.
    """

    def __init__(self, output_dir, scan_timestamp, timeout=15, viewport=None):
        """
        Initialize the screenshot capture handler.

        Args:
            output_dir (str): Base output directory for scan reports
            scan_timestamp (str): Timestamp string for this scan (format: YYYYMMDD_HHMMSS)
            timeout (int): Timeout in seconds for page load and screenshot (default: 15)
            viewport (dict): Viewport size dict with 'width' and 'height' keys
                           (default: {'width': 1280, 'height': 720})
        """
        self.output_dir = output_dir
        self.scan_timestamp = scan_timestamp
        self.timeout = timeout * 1000  # Convert to milliseconds for Playwright
        self.viewport = viewport or {'width': 1280, 'height': 720}

        self.playwright = None
        self.browser = None
        self.screenshot_dir = None

        # Set up logging
        self.logger = logging.getLogger('SneakyScanner.Screenshot')

    def _get_screenshot_dir(self):
        """
        Create and return the screenshots subdirectory for this scan.

        Returns:
            Path: Path object for the screenshots directory
        """
        if self.screenshot_dir is None:
            dir_name = f"scan_report_{self.scan_timestamp}_screenshots"
            self.screenshot_dir = Path(self.output_dir) / dir_name
            self.screenshot_dir.mkdir(parents=True, exist_ok=True)
            self.logger.info(f"Created screenshot directory: {self.screenshot_dir}")

        return self.screenshot_dir

    def _generate_filename(self, ip, port):
        """
        Generate a filename for the screenshot.

        Args:
            ip (str): IP address of the service
            port (int): Port number of the service

        Returns:
            str: Filename in format: {ip}_{port}.png
        """
        # Replace dots in IP with underscores for filesystem compatibility
        safe_ip = ip.replace('.', '_')
        return f"{safe_ip}_{port}.png"

    def _launch_browser(self):
        """
        Launch Playwright and Chromium browser in headless mode.

        Returns:
            bool: True if browser launched successfully, False otherwise
        """
        if self.browser is not None:
            return True  # Already launched

        try:
            self.logger.info("Launching Chromium browser...")
            self.playwright = sync_playwright().start()
            self.browser = self.playwright.chromium.launch(
                headless=True,
                args=[
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--disable-dev-shm-usage',
                    '--disable-gpu',
                ]
            )
            self.logger.info("Chromium browser launched successfully")
            return True

        except Exception as e:
            self.logger.error(f"Failed to launch browser: {e}")
            return False

    def _close_browser(self):
        """
        Close the browser and cleanup Playwright resources.
        """
        if self.browser:
            try:
                self.browser.close()
                self.logger.info("Browser closed")
            except Exception as e:
                self.logger.warning(f"Error closing browser: {e}")
            finally:
                self.browser = None

        if self.playwright:
            try:
                self.playwright.stop()
            except Exception as e:
                self.logger.warning(f"Error stopping playwright: {e}")
            finally:
                self.playwright = None

    def capture(self, ip, port, protocol):
        """
        Capture a screenshot of a web service.

        Args:
            ip (str): IP address of the service
            port (int): Port number of the service
            protocol (str): Protocol to use ('http' or 'https')

        Returns:
            str: Relative path to the screenshot file, or None if capture failed
        """
        # Validate protocol
        if protocol not in ['http', 'https']:
            self.logger.warning(f"Invalid protocol '{protocol}' for {ip}:{port}")
            return None

        # Launch browser if not already running
        if not self._launch_browser():
            return None

        # Build URL
        url = f"{protocol}://{ip}:{port}"

        # Generate screenshot filename
        filename = self._generate_filename(ip, port)
        screenshot_dir = self._get_screenshot_dir()
        screenshot_path = screenshot_dir / filename

        try:
            self.logger.info(f"Capturing screenshot: {url}")

            # Create new browser context with viewport and SSL settings
            context = self.browser.new_context(
                viewport=self.viewport,
                ignore_https_errors=True,  # Handle self-signed certs
                user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            )

            # Create new page
            page = context.new_page()

            # Set default timeout
            page.set_default_timeout(self.timeout)

            # Navigate to URL
            page.goto(url, wait_until='networkidle', timeout=self.timeout)

            # Take screenshot (viewport only)
            page.screenshot(path=str(screenshot_path), type='png')

            # Close page and context
            page.close()
            context.close()

            self.logger.info(f"Screenshot saved: {screenshot_path}")

            # Return relative path (relative to output directory)
            relative_path = f"{screenshot_dir.name}/{filename}"
            return relative_path

        except PlaywrightTimeout:
            self.logger.warning(f"Timeout capturing screenshot for {url}")
            return None

        except Exception as e:
            self.logger.warning(f"Failed to capture screenshot for {url}: {e}")
            return None

    def __enter__(self):
        """Context manager entry."""
        self._launch_browser()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - cleanup browser resources."""
        self._close_browser()
        return False  # Don't suppress exceptions