"""Screenshot capture for visual regression.

Supports:
- Playwright (recommended)
- Selenium WebDriver
- Raw screenshot bytes

Integrates with Hikaflow session replay timeline.
"""
from __future__ import annotations

import hashlib
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional

if TYPE_CHECKING:
    from playwright.async_api import Page as AsyncPlaywrightPage
    from playwright.sync_api import Page as PlaywrightPage
    from selenium.webdriver.remote.webdriver import WebDriver

logger = logging.getLogger(__name__)


@dataclass
class VisualSnapshot:
    """A captured visual snapshot."""

    name: str
    screenshot: bytes  # PNG bytes
    timestamp: datetime
    url: Optional[str] = None
    viewport_size: Optional[Dict[str, int]] = None
    device_scale_factor: float = 1.0
    full_page: bool = False
    element_selector: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def content_hash(self) -> str:
        """SHA256 hash of screenshot content."""
        return hashlib.sha256(self.screenshot).hexdigest()[:16]

    @property
    def filename(self) -> str:
        """Generate filename for this snapshot."""
        safe_name = self.name.replace("/", "_").replace(" ", "_")
        return f"{safe_name}_{self.content_hash}.png"

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary (without screenshot bytes)."""
        return {
            "name": self.name,
            "timestamp": self.timestamp.isoformat(),
            "url": self.url,
            "viewport_size": self.viewport_size,
            "device_scale_factor": self.device_scale_factor,
            "full_page": self.full_page,
            "element_selector": self.element_selector,
            "content_hash": self.content_hash,
            "metadata": self.metadata,
        }


class VisualCapture:
    """Capture screenshots for visual regression testing.

    Supports multiple capture backends:
    - Playwright (sync and async)
    - Selenium WebDriver
    - Custom callback

    Usage with Playwright:
        from playwright.sync_api import sync_playwright

        with sync_playwright() as p:
            browser = p.chromium.launch()
            page = browser.new_page()
            page.goto("https://example.com")

            capture = VisualCapture.from_playwright(page)
            snapshot = capture.take("homepage")

    Usage with Selenium:
        from selenium import webdriver

        driver = webdriver.Chrome()
        driver.get("https://example.com")

        capture = VisualCapture.from_selenium(driver)
        snapshot = capture.take("homepage")
    """

    def __init__(
        self,
        capture_fn: Callable[[], bytes],
        url_fn: Optional[Callable[[], str]] = None,
        viewport_fn: Optional[Callable[[], Dict[str, int]]] = None,
        device_scale_factor: float = 1.0,
        element_capture_fn: Optional[Callable[[str], bytes]] = None,
    ):
        """Initialize capture with custom functions.

        Args:
            capture_fn: Function that returns screenshot PNG bytes
            url_fn: Function that returns current URL
            viewport_fn: Function that returns viewport size
            device_scale_factor: Device pixel ratio
            element_capture_fn: Optional function that takes a CSS selector
                and returns screenshot PNG bytes of that element (P7-F9).
                When provided, take_element() and take(element_selector=...)
                will capture only the targeted element.
        """
        self._capture_fn = capture_fn
        self._url_fn = url_fn
        self._viewport_fn = viewport_fn
        self._device_scale_factor = device_scale_factor
        self._element_capture_fn = element_capture_fn
        self._snapshots: List[VisualSnapshot] = []

    @classmethod
    def from_playwright(
        cls,
        page: PlaywrightPage,
        device_scale_factor: float = 1.0,
    ) -> VisualCapture:
        """Create capture from Playwright page.

        Args:
            page: Playwright Page object
            device_scale_factor: Device pixel ratio
        """
        def capture_fn() -> bytes:
            return page.screenshot(type="png")

        def url_fn() -> str:
            return page.url

        def viewport_fn() -> Dict[str, int]:
            size = page.viewport_size
            return {"width": size["width"], "height": size["height"]} if size else {}

        def element_capture_fn(selector: str) -> bytes:
            """Capture screenshot of a specific element via locator (P7-F9)."""
            return page.locator(selector).screenshot(type="png")

        return cls(
            capture_fn=capture_fn,
            url_fn=url_fn,
            viewport_fn=viewport_fn,
            device_scale_factor=device_scale_factor,
            element_capture_fn=element_capture_fn,
        )

    @classmethod
    def from_playwright_async(
        cls,
        page: AsyncPlaywrightPage,
        device_scale_factor: float = 1.0,
    ) -> AsyncVisualCapture:
        """Create async capture from Playwright page.

        Args:
            page: Playwright async Page object
            device_scale_factor: Device pixel ratio
        """
        return AsyncVisualCapture(page, device_scale_factor)

    @classmethod
    def from_selenium(
        cls,
        driver: WebDriver,
        device_scale_factor: float = 1.0,
    ) -> VisualCapture:
        """Create capture from Selenium WebDriver.

        Args:
            driver: Selenium WebDriver
            device_scale_factor: Device pixel ratio
        """
        def capture_fn() -> bytes:
            return driver.get_screenshot_as_png()

        def url_fn() -> str:
            return driver.current_url

        def viewport_fn() -> Dict[str, int]:
            size = driver.get_window_size()
            return {"width": size["width"], "height": size["height"]}

        return cls(
            capture_fn=capture_fn,
            url_fn=url_fn,
            viewport_fn=viewport_fn,
            device_scale_factor=device_scale_factor,
        )

    def take(
        self,
        name: str,
        full_page: bool = False,
        element_selector: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> VisualSnapshot:
        """Take a visual snapshot.

        Args:
            name: Unique name for this snapshot
            full_page: Capture full scrollable page
            element_selector: CSS selector to capture specific element
            metadata: Additional metadata to attach

        Returns:
            VisualSnapshot
        """
        # Capture screenshot -- use element-level capture when available (P7-F9)
        if element_selector and self._element_capture_fn is not None:
            try:
                screenshot = self._element_capture_fn(element_selector)
            except Exception as e:
                logger.warning(
                    f"Element capture failed for '{element_selector}', "
                    f"falling back to full page: {e}"
                )
                screenshot = self._capture_fn()
        else:
            screenshot = self._capture_fn()

        # Get context
        url = self._url_fn() if self._url_fn else None
        viewport = self._viewport_fn() if self._viewport_fn else None

        snapshot = VisualSnapshot(
            name=name,
            screenshot=screenshot,
            timestamp=datetime.now(timezone.utc),
            url=url,
            viewport_size=viewport,
            device_scale_factor=self._device_scale_factor,
            full_page=full_page,
            element_selector=element_selector,
            metadata=metadata or {},
        )

        self._snapshots.append(snapshot)
        logger.debug(f"Captured snapshot: {name} ({len(screenshot)} bytes)")

        return snapshot

    def take_element(
        self,
        name: str,
        selector: str,
        padding: int = 0,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> VisualSnapshot:
        """Take snapshot of a specific element.

        Note: For Playwright, this uses the page-level screenshot.
        For element-specific screenshots with Playwright, use the
        locator.screenshot() method directly.

        Args:
            name: Snapshot name
            selector: CSS selector for element
            padding: Pixels of padding around element
            metadata: Additional metadata
        """
        return self.take(
            name=name,
            element_selector=selector,
            metadata={"padding": padding, **(metadata or {})},
        )

    @property
    def snapshots(self) -> List[VisualSnapshot]:
        """Get all captured snapshots."""
        return self._snapshots.copy()

    def clear(self) -> None:
        """Clear captured snapshots."""
        self._snapshots.clear()


class AsyncVisualCapture:
    """Async version of VisualCapture for Playwright async API."""

    def __init__(
        self,
        page: AsyncPlaywrightPage,
        device_scale_factor: float = 1.0,
    ):
        self._page = page
        self._device_scale_factor = device_scale_factor
        self._snapshots: List[VisualSnapshot] = []

    async def take(
        self,
        name: str,
        full_page: bool = False,
        element_selector: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> VisualSnapshot:
        """Take a visual snapshot asynchronously."""
        # Capture screenshot
        if element_selector:
            element = self._page.locator(element_selector)
            screenshot = await element.screenshot(type="png")
        else:
            screenshot = await self._page.screenshot(type="png", full_page=full_page)

        # Get context
        url = self._page.url
        viewport = self._page.viewport_size
        viewport_dict = {"width": viewport["width"], "height": viewport["height"]} if viewport else None

        snapshot = VisualSnapshot(
            name=name,
            screenshot=screenshot,
            timestamp=datetime.now(timezone.utc),
            url=url,
            viewport_size=viewport_dict,
            device_scale_factor=self._device_scale_factor,
            full_page=full_page,
            element_selector=element_selector,
            metadata=metadata or {},
        )

        self._snapshots.append(snapshot)
        logger.debug(f"Captured async snapshot: {name} ({len(screenshot)} bytes)")

        return snapshot

    async def take_element(
        self,
        name: str,
        selector: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> VisualSnapshot:
        """Take snapshot of a specific element."""
        return await self.take(
            name=name,
            element_selector=selector,
            metadata=metadata,
        )

    @property
    def snapshots(self) -> List[VisualSnapshot]:
        """Get all captured snapshots."""
        return self._snapshots.copy()

    def clear(self) -> None:
        """Clear captured snapshots."""
        self._snapshots.clear()


def capture_from_file(image_path: Path) -> VisualSnapshot:
    """Create snapshot from existing image file.

    Useful for testing or when screenshots are captured externally.

    Args:
        image_path: Path to PNG image

    Returns:
        VisualSnapshot
    """
    with open(image_path, "rb") as f:
        screenshot = f.read()

    return VisualSnapshot(
        name=image_path.stem,
        screenshot=screenshot,
        timestamp=datetime.now(timezone.utc),
        metadata={"source": "file", "path": str(image_path)},
    )
