#!/usr/bin/env python3
"""
Vision Worker for UVS
=====================
Dedicated worker for high-frequency (16 FPS) vision capture and streaming.
Bridges BrowserController with Gemini Multimodal Live API.

VERIFICATION_STAMP
Story: UVS-H10, UVS-H15, UVS-H28
- UVS-H10: Proper task tracking and cleanup
- UVS-H15: Vision task cancellation handling
- UVS-H28: Backpressure with bounded frame buffer
"""

import asyncio
import logging
from typing import Optional, Callable, Any
from collections import deque

logger = logging.getLogger(__name__)

# UVS-H28: Maximum pending frames to prevent memory bloat
MAX_PENDING_FRAMES = 5


class VisionWorker:
    """
    Dedicated worker for high-frequency (16 FPS) vision capture and streaming.
    Bridges BrowserController with Gemini Multimodal Live API.
    """

    def __init__(self, browser_controller: Any, fps: float = 16.0):
        self.controller = browser_controller
        self.fps = fps
        self.frame_interval = 1.0 / fps
        self.is_running = False
        self._task: Optional[asyncio.Task] = None
        self.on_frame: Optional[Callable[[bytes], Any]] = None
        self.quality = 50  # JPEG Compression Quality
        self.scale = 1.0   # Optional downscaling

        # UVS-H28: Bounded frame buffer for backpressure
        self._pending_frames: deque = deque(maxlen=MAX_PENDING_FRAMES)
        self._frames_dropped = 0

    async def start(self):
        """Starts the vision capture loop."""
        if self.is_running:
            return

        if not self.controller._initialized:
            logger.warning("VisionWorker: BrowserController not initialized.")
            return

        self.is_running = True
        self._frames_dropped = 0

        # UVS-H10: Create task with done callback for exception logging
        self._task = asyncio.create_task(self._loop())
        self._task.add_done_callback(self._task_done_callback)
        logger.info(f"VisionWorker started at {self.fps} FPS")

    def _task_done_callback(self, task: asyncio.Task):
        """UVS-H10: Log any exceptions from the vision task."""
        if task.cancelled():
            logger.debug("VisionWorker task cancelled")
        elif task.exception():
            logger.error(f"VisionWorker task exception: {task.exception()}")

    async def stop(self):
        """
        Stops the vision capture loop.

        UVS-H15: Proper cancellation handling with exception logging.
        """
        self.is_running = False

        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                logger.debug("VisionWorker task cancelled during stop")
            except Exception as e:
                # UVS-H15: Log any other exceptions during cleanup
                logger.warning(f"VisionWorker stop exception: {e}")
            finally:
                self._task = None

        # Clear frame buffer
        self._pending_frames.clear()

        if self._frames_dropped > 0:
            logger.info(f"VisionWorker stopped. Frames dropped due to backpressure: {self._frames_dropped}")
        else:
            logger.info("VisionWorker stopped")

    async def _loop(self):
        """Main capture loop with backpressure handling."""
        last_frame_time = 0

        while self.is_running:
            now = asyncio.get_event_loop().time()
            if now - last_frame_time < self.frame_interval:
                await asyncio.sleep(0.01)
                continue

            try:
                # Optimized screenshot from controller
                screenshot_bytes = await self._capture_frame()

                if screenshot_bytes and self.on_frame:
                    # UVS-H28: Backpressure - if pending frames > limit, drop oldest
                    if len(self._pending_frames) >= MAX_PENDING_FRAMES:
                        self._frames_dropped += 1
                        logger.debug("Frame dropped due to backpressure")

                    # Add to pending queue (deque maxlen handles overflow)
                    self._pending_frames.append(screenshot_bytes)

                    # Process the callback
                    if asyncio.iscoroutinefunction(self.on_frame):
                        await self.on_frame(screenshot_bytes)
                    else:
                        self.on_frame(screenshot_bytes)

                last_frame_time = now

            except asyncio.CancelledError:
                # UVS-H15: Handle cancellation cleanly
                logger.debug("VisionWorker loop cancelled")
                break
            except Exception as e:
                logger.error(f"VisionWorker loop error: {e}")
                await asyncio.sleep(0.5)

    async def _capture_frame(self) -> Optional[bytes]:
        """Captures and processes a single frame."""
        # Note: We rely on the controller's active backend (Playwright/etc)
        backend = self.controller._active_backend
        if not backend or not hasattr(backend, "_page") or not backend._page:
            return None

        try:
            # Direct Playwright screenshot with compression
            # type="jpeg" and quality are supported by Playwright
            screenshot_bytes = await backend._page.screenshot(
                type="jpeg",
                quality=self.quality,
                scale="css"  # Viewport scaling
            )

            return screenshot_bytes

        except Exception as e:
            logger.debug(f"Frame capture failed: {e}")
            return None

    def set_fps(self, fps: float):
        """Update target FPS."""
        self.fps = max(1.0, min(fps, 60.0))  # Clamp to 1-60 FPS
        self.frame_interval = 1.0 / self.fps

    def get_stats(self) -> dict:
        """Get vision worker statistics."""
        return {
            "is_running": self.is_running,
            "fps": self.fps,
            "quality": self.quality,
            "pending_frames": len(self._pending_frames),
            "frames_dropped": self._frames_dropped
        }