#!/usr/bin/env python3
# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
# One which is durable, easy to repair, & distributed like tap water
# for machine learning intelligence.
#
# The permacomputer is community-owned infrastructure optimized around
# four values:
#
#   TRUTH      First principles, math & science, open source code freely distributed
#   FREEDOM    Voluntary partnerships, freedom from tyranny & corporate control
#   HARMONY    Minimal waste, self-renewing systems with diverse thriving connections
#   LOVE       Be yourself without hurting others, cooperation through natural law
#
# This software contributes to that vision by archiving the web, preserving digital knowledge before it disappears.
# Code is seeds to sprout on any abandoned technology.

"""
Screenshot capture module for neopig.

Wraps uri2png for async-compatible page screenshots.
Screenshots are stored in vault with MD5 hash like other media.

Supported engines (in order of preference):
  - wkhtmltoimage: Fast, lightweight, uses Qt WebKit. Install: apt install wkhtmltopdf
  - cutycapt: Fast, lightweight, uses Qt WebKit. Install: apt install cutycapt
  - playwright-webkit: WebKit via Playwright (lighter than Chromium)
  - playwright-firefox: Firefox via Playwright
  - playwright-chromium: Chromium via Playwright (heaviest, but most compatible)
  - selenium-*: Various Selenium drivers

The module auto-detects available engines and picks the lightest one,
or you can specify an engine explicitly.
"""

import asyncio
import hashlib
import io
import logging
import shutil
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, List, Dict, Any
import tempfile

from PIL import Image

# Suppress PIL decompression bomb warnings for large screenshots
warnings.filterwarnings('ignore', category=Image.DecompressionBombWarning)

logger = logging.getLogger(__name__)

# Engine preference order - lightest/fastest first
ENGINE_PREFERENCE = [
    'cutycapt',         # Native Qt WebKit - fast, supports full-page screenshots
    'wkhtmltoimage',    # Native Qt WebKit - fast, viewport-only
    'playwright-webkit', # WebKit via Playwright - lighter than Chromium
    'playwright-firefox',
    'playwright-chromium',
    'playwright',       # Default Playwright (Chromium)
    'selenium-chrome',
    'selenium-firefox',
    'selenium',
]

# Native tools that don't require browser downloads
NATIVE_ENGINES = {'wkhtmltoimage', 'cutycapt'}


@dataclass
class ScreenshotConfig:
    """Screenshot capture configuration."""
    enabled: bool = True
    width: int = 1024
    height: int = 768
    delay: int = 1000  # ms after DOM load
    timeout: int = 30000  # ms total timeout
    user_agent: Optional[str] = None
    engine: Optional[str] = None  # None = auto-detect best available
    full_page: bool = True
    format: str = 'jpeg'  # 'jpeg' or 'png'
    quality: int = 93  # JPEG quality (1-100)


class ScreenshotCapture:
    """
    Async screenshot capture using uri2png.

    Supports multiple backends with automatic selection of the lightest
    available engine. Native tools (wkhtmltoimage, cutycapt) are preferred
    over browser-based solutions.

    Usage:
        config = ScreenshotConfig(enabled=True, engine='wkhtmltoimage')
        capture = ScreenshotCapture(config)
        result = await capture.capture('https://example.com')
    """

    def __init__(self, config: ScreenshotConfig = None):
        self.config = config or ScreenshotConfig()
        self._engine = None
        self._engine_name = None
        self._available_engines: Optional[List[Dict[str, str]]] = None
        self._initialized = False

    async def _get_available_engines(self) -> List[Dict[str, str]]:
        """Get list of available screenshot engines."""
        if self._available_engines is not None:
            return self._available_engines

        def _check():
            try:
                from uri2png import get_available_engines
                return get_available_engines()
            except ImportError:
                return []

        self._available_engines = await asyncio.to_thread(_check)
        return self._available_engines

    async def _select_engine(self) -> Optional[str]:
        """Select the best available engine based on preference order."""
        available = await self._get_available_engines()
        available_names = {e['name'] for e in available}

        # If user specified an engine, try to use it
        if self.config.engine:
            if self.config.engine in available_names:
                return self.config.engine
            else:
                logger.warning(f"Requested engine '{self.config.engine}' not available")
                logger.info(f"Available engines: {', '.join(available_names)}")

        # Check native tools first (they're fast and don't need browser downloads)
        for engine in ENGINE_PREFERENCE:
            if engine in available_names:
                # For native engines, verify the binary exists
                if engine in NATIVE_ENGINES:
                    binary = 'wkhtmltoimage' if engine == 'wkhtmltoimage' else 'cutycapt'
                    if shutil.which(binary):
                        return engine
                    else:
                        logger.debug(f"Engine {engine} listed but binary not found")
                        continue
                return engine

        return None

    async def initialize(self) -> bool:
        """Initialize the screenshot engine."""
        if self._initialized:
            return self._engine is not None

        engine_name = await self._select_engine()
        if not engine_name:
            logger.warning("No screenshot engine available")
            logger.info("Install one of: wkhtmltopdf, cutycapt, or playwright")
            self._initialized = True
            return False

        def _create_engine():
            try:
                from uri2png import create_engine
                # Pass options as kwargs to create_engine
                return create_engine(
                    engine_name,
                    width=self.config.width,
                    height=self.config.height,
                    delay=self.config.delay,
                    timeout=self.config.timeout,
                    full_page=self.config.full_page,
                    user_agent=self.config.user_agent,
                )
            except Exception as e:
                logger.warning(f"Failed to create engine '{engine_name}': {e}")
                return None

        self._engine = await asyncio.to_thread(_create_engine)
        self._engine_name = engine_name
        self._initialized = True

        if self._engine:
            logger.info(f"Screenshot engine: {engine_name}")
            return True
        return False

    async def is_available(self) -> bool:
        """Check if screenshots are available."""
        if not self._initialized:
            await self.initialize()
        return self._engine is not None

    def get_engine_name(self) -> Optional[str]:
        """Get the name of the active engine."""
        return self._engine_name

    async def list_engines(self) -> List[Dict[str, str]]:
        """List all available screenshot engines."""
        return await self._get_available_engines()

    @staticmethod
    def calculate_delay_for_content(content_length: int) -> int:
        """
        Calculate optimal screenshot delay based on content size.

        Longer pages need more time for JavaScript to render all content.

        Args:
            content_length: Length of raw HTML content in bytes

        Returns:
            Delay in milliseconds (min 3000ms, max 60000ms)
        """
        # Base delay of 3 seconds
        base_delay = 3000

        # Add 3 seconds per 50KB of content, up to a max of 57 seconds
        additional_delay = min((content_length // 50000) * 3000, 57000)

        return base_delay + additional_delay

    async def capture(self, uri: str, content_length: int = 0) -> Optional[dict]:
        """
        Capture screenshot of a URI.

        Args:
            uri: URL to capture
            content_length: Optional content length for dynamic delay calculation.
                           If provided and engine supports it, delay will be adjusted.

        Returns:
            Dict with 'data', 'md5_hash', 'mime_type', 'engine' or None on failure
        """
        if not self.config.enabled:
            return None

        if not await self.is_available():
            return None

        # Calculate dynamic delay if content_length provided and we're using a native engine
        effective_engine = self._engine
        if content_length > 0 and self._engine_name in NATIVE_ENGINES:
            dynamic_delay = self.calculate_delay_for_content(content_length)
            if dynamic_delay > self.config.delay:
                logger.debug(f"Using dynamic delay {dynamic_delay}ms for {content_length} bytes")
                # Create a one-off engine with the adjusted delay
                try:
                    from uri2png import create_engine
                    effective_engine = await asyncio.to_thread(
                        lambda: create_engine(
                            self._engine_name,
                            width=self.config.width,
                            height=self.config.height,
                            delay=dynamic_delay,
                            timeout=self.config.timeout,
                            full_page=self.config.full_page,
                            user_agent=self.config.user_agent,
                        )
                    )
                except Exception as e:
                    logger.debug(f"Could not create dynamic engine: {e}")
                    effective_engine = self._engine

        try:
            # Create temp file for output
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
                output_path = f.name

            try:
                # Capture - API is capture(url, output_path)
                capture_coro = effective_engine.capture(uri, output_path)

                # All uri2png engines return coroutines
                result = await asyncio.wait_for(
                    capture_coro,
                    timeout=self.config.timeout / 1000 + 5
                )

                if not result.success:
                    # Extract meaningful error from wkhtmltoimage output
                    raw_error = result.error or 'unknown error'
                    # Filter out wkhtmltoimage progress output (progress bars, loading messages)
                    lines = raw_error.split('\n')
                    # Skip lines that are progress bars or loading messages
                    meaningful = [l.strip() for l in lines if l.strip()
                                  and not l.strip().startswith('[')  # Progress bars like [>
                                  and not l.strip().startswith('Loading')
                                  and '%' not in l]  # Percentage indicators
                    # Look for actual error/fail lines first
                    error_lines = [l for l in meaningful if 'error' in l.lower() or 'fail' in l.lower()]
                    if error_lines:
                        error_msg = error_lines[0][:60]
                    elif meaningful:
                        error_msg = meaningful[-1][:60]
                    else:
                        error_msg = 'capture failed'
                    logger.warning(f"Screenshot failed: {uri} - {error_msg}")
                    return None

                # Read bytes from output file
                png_data = Path(output_path).read_bytes()
                if not png_data:
                    logger.warning(f"Screenshot empty: {uri}")
                    return None

                # Convert to JPEG if configured
                if self.config.format == 'jpeg':
                    img = Image.open(io.BytesIO(png_data))

                    # Convert RGBA to RGB (JPEG doesn't support alpha)
                    if img.mode in ('RGBA', 'LA', 'P'):
                        img = img.convert('RGB')

                    # Split oversized images (JPEG max is 65500px)
                    max_dim = 65000
                    if img.height > max_dim:
                        logger.debug(f"Splitting oversized image: {img.height}px into chunks")
                        results = []
                        chunk_idx = 0
                        y = 0
                        while y < img.height:
                            chunk_height = min(max_dim, img.height - y)
                            chunk = img.crop((0, y, img.width, y + chunk_height))
                            output = io.BytesIO()
                            chunk.save(output, format='JPEG', quality=self.config.quality, optimize=True)
                            chunk_data = output.getvalue()
                            chunk_hash = hashlib.md5(chunk_data).hexdigest()
                            results.append({
                                'data': chunk_data,
                                'md5_hash': chunk_hash,
                                'mime_type': 'image/jpeg',
                                'size': len(chunk_data),
                                'source_uri': uri,
                                'engine': self._engine_name,
                                'format': 'jpg',
                                'chunk': chunk_idx,
                                'total_chunks': (img.height + max_dim - 1) // max_dim,
                            })
                            y += max_dim
                            chunk_idx += 1
                        return results  # Return list for oversized images

                    output = io.BytesIO()
                    img.save(output, format='JPEG', quality=self.config.quality, optimize=True)
                    data = output.getvalue()
                    mime_type = 'image/jpeg'
                    ext = 'jpg'
                else:
                    data = png_data
                    mime_type = 'image/png'
                    ext = 'png'

                md5_hash = hashlib.md5(data).hexdigest()
                logger.debug(f"Screenshot captured ({self._engine_name}, {ext}): {uri} -> {md5_hash} ({len(data)//1024}KB)")

                return {
                    'data': data,
                    'md5_hash': md5_hash,
                    'mime_type': mime_type,
                    'size': len(data),
                    'source_uri': uri,
                    'engine': self._engine_name,
                    'format': ext,
                }
            finally:
                # Clean up temp file
                try:
                    Path(output_path).unlink(missing_ok=True)
                except Exception:
                    pass

        except asyncio.TimeoutError:
            logger.warning(f"Screenshot timeout: {uri}")
            return None
        except Exception as e:
            logger.warning(f"Screenshot error for {uri}: {e}")
            return None

    async def capture_to_file(self, uri: str, output_path: str) -> bool:
        """
        Capture screenshot directly to a file.

        Returns:
            True on success, False on failure
        """
        result = await self.capture(uri)
        if result is None:
            return False

        try:
            await asyncio.to_thread(Path(output_path).write_bytes, result['data'])
            return True
        except Exception as e:
            logger.warning(f"Failed to write screenshot: {e}")
            return False

    async def cleanup(self):
        """Cleanup engine resources."""
        if self._engine:
            try:
                def _cleanup():
                    if hasattr(self._engine, 'cleanup'):
                        self._engine.cleanup()
                await asyncio.to_thread(_cleanup)
            except Exception as e:
                logger.debug(f"Engine cleanup error: {e}")


async def list_available_engines() -> List[Dict[str, str]]:
    """List all available screenshot engines (utility function)."""
    capture = ScreenshotCapture()
    return await capture.list_engines()


async def get_best_engine() -> Optional[str]:
    """Get the best available engine name (utility function)."""
    capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
    await capture.initialize()
    return capture.get_engine_name()
