"""Hybrid Perception Engine — combines screenshot + UI tree + OCR for rich state understanding.

This provides a more accurate device state representation than screenshot alone,
combining:
1. Screenshot (visual/pixel-level information)
2. UI Tree (semantic structure via Accessibility/UIAutomator)
3. OCR (text recognition for non-accessible text in images/games)

The combined state is passed to the VLM for better action planning.
"""

from __future__ import annotations

import base64
import logging
from dataclasses import dataclass, field
from typing import Any

from ..runtime.base import DeviceRuntime

logger = logging.getLogger(__name__)


@dataclass
class PerceptionResult:
    """Rich device state from hybrid perception."""
    screenshot_b64: str
    ui_elements: list[UIElement]
    ocr_texts: list[OCRText]
    screen_width: int
    screen_height: int

    def to_prompt(self, max_elements: int = 80) -> str:
        """Convert to a structured text prompt for the VLM."""
        lines = [f"Screen: {self.screen_width}x{self.screen_height}"]
        lines.append("")

        if self.ui_elements:
            lines.append("=== UI Elements (from accessibility tree) ===")
            for i, elem in enumerate(self.ui_elements[:max_elements]):
                parts = [f"[{i}] {elem.class_name}"]
                if elem.text:
                    parts.append(f'text="{elem.text}"')
                if elem.content_desc:
                    parts.append(f'desc="{elem.content_desc}"')
                if elem.bounds:
                    parts.append(f"bounds={elem.bounds}")
                flags = []
                if elem.clickable:
                    flags.append("clickable")
                if elem.scrollable:
                    flags.append("scrollable")
                if elem.focused:
                    flags.append("focused")
                if elem.selected:
                    flags.append("selected")
                if flags:
                    parts.append(f"({', '.join(flags)})")
                lines.append(" ".join(parts))
            if len(self.ui_elements) > max_elements:
                lines.append(f"... and {len(self.ui_elements) - max_elements} more elements")

        if self.ocr_texts:
            lines.append("")
            lines.append("=== OCR-detected text (not in UI tree) ===")
            for ocr in self.ocr_texts[:20]:
                lines.append(f'  "{ocr.text}" at ({ocr.x}, {ocr.y}) size={ocr.w}x{ocr.h} conf={ocr.confidence:.0%}')

        return "\n".join(lines)


@dataclass
class UIElement:
    class_name: str
    text: str = ""
    content_desc: str = ""
    resource_id: str = ""
    bounds: str = ""
    clickable: bool = False
    scrollable: bool = False
    focusable: bool = False
    focused: bool = False
    selected: bool = False
    center_x: int = 0
    center_y: int = 0


@dataclass
class OCRText:
    text: str
    x: int
    y: int
    w: int
    h: int
    confidence: float = 0.0


class HybridPerceptionEngine:
    """Combines multiple perception modalities for rich state understanding."""

    def __init__(
        self,
        runtime: DeviceRuntime,
        enable_ocr: bool = False,
    ) -> None:
        self._runtime = runtime
        self._enable_ocr = enable_ocr
        self._ocr_reader: Any = None

    async def perceive(self) -> PerceptionResult:
        """Capture the full device state."""
        screenshot_bytes = await self._runtime.screenshot()
        screenshot_b64 = base64.b64encode(screenshot_bytes).decode()

        info = await self._runtime.get_info()

        ui_tree_raw = await self._runtime.get_ui_tree()
        ui_elements = self._parse_ui_tree(ui_tree_raw)

        ocr_texts: list[OCRText] = []
        if self._enable_ocr and screenshot_bytes:
            ocr_texts = await self._run_ocr(screenshot_bytes)

        return PerceptionResult(
            screenshot_b64=screenshot_b64,
            ui_elements=ui_elements,
            ocr_texts=ocr_texts,
            screen_width=info.width,
            screen_height=info.height,
        )

    def _parse_ui_tree(self, nodes: list[dict]) -> list[UIElement]:
        """Convert raw UI tree nodes into structured UIElement objects."""
        elements: list[UIElement] = []
        for node in nodes:
            raw_bounds = node.get("bounds", "")
            cx, cy = 0, 0
            bounds_display = ""

            if isinstance(raw_bounds, dict):
                left = raw_bounds.get("left", 0)
                top = raw_bounds.get("top", 0)
                right = raw_bounds.get("right", 0)
                bottom = raw_bounds.get("bottom", 0)
                cx = (left + right) // 2
                cy = (top + bottom) // 2
                bounds_display = f"[{left},{top}][{right},{bottom}]"
            elif isinstance(raw_bounds, str) and raw_bounds:
                cx, cy = self._parse_bounds_center(raw_bounds)
                bounds_display = raw_bounds

            elements.append(UIElement(
                class_name=node.get("class", "View"),
                text=node.get("text", ""),
                content_desc=node.get("content-desc", ""),
                resource_id=node.get("resource-id", ""),
                bounds=bounds_display,
                clickable=node.get("clickable", False) in (True, "true"),
                scrollable=node.get("scrollable", False) in (True, "true"),
                focusable=node.get("focusable", False) in (True, "true"),
                focused=node.get("focused", False) in (True, "true"),
                selected=node.get("selected", False) in (True, "true"),
                center_x=cx,
                center_y=cy,
            ))
        return elements

    @staticmethod
    def _parse_bounds_center(bounds_str: str) -> tuple[int, int]:
        """Parse Android bounds format [x1,y1][x2,y2] to center point."""
        try:
            parts = bounds_str.replace("][", ",").strip("[]").split(",")
            if len(parts) == 4:
                x1, y1, x2, y2 = int(parts[0]), int(parts[1]), int(parts[2]), int(parts[3])
                return (x1 + x2) // 2, (y1 + y2) // 2
        except (ValueError, IndexError):
            pass
        return 0, 0

    def _get_ocr_reader(self) -> Any:
        """Return a cached easyocr.Reader instance, creating one on first call."""
        if self._ocr_reader is not None:
            return self._ocr_reader
        import easyocr
        try:
            self._ocr_reader = easyocr.Reader(["en", "ch_sim"], gpu=True)
        except RuntimeError:
            logger.info("GPU not available for OCR, falling back to CPU")
            self._ocr_reader = easyocr.Reader(["en", "ch_sim"], gpu=False)
        return self._ocr_reader

    async def _run_ocr(self, screenshot_bytes: bytes) -> list[OCRText]:
        """Run OCR on the screenshot to detect text not in the UI tree."""
        try:
            import easyocr  # noqa: F401
        except ImportError:
            logger.debug("easyocr not installed, skipping OCR")
            return []

        import io
        from PIL import Image
        import numpy as np

        try:
            img = Image.open(io.BytesIO(screenshot_bytes))
            arr = np.array(img)

            reader = self._get_ocr_reader()
            results = reader.readtext(arr)

            texts = []
            for bbox, text, conf in results:
                if conf < 0.3:
                    continue
                x1, y1 = int(bbox[0][0]), int(bbox[0][1])
                x2, y2 = int(bbox[2][0]), int(bbox[2][1])
                texts.append(OCRText(
                    text=text,
                    x=x1, y=y1,
                    w=x2 - x1, h=y2 - y1,
                    confidence=conf,
                ))
            return texts
        except Exception as e:
            logger.warning("OCR failed: %s", e)
            return []
