from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from typing import Any

from core.models import Scene
from perception.ocr import PaddleOCREngine, run_ocr
from perception.scene_builder import build_scene as build_fused_scene
from perception.scene_builder import infer_png_size
from perception.ui_parser import parse_ui_tree

logger = logging.getLogger(__name__)


class PerceptionProvider(ABC):
    @abstractmethod
    def build_scene(self, screenshot: bytes, tree: Any) -> Scene:
        """Build a normalized scene from raw device observations."""


class DefaultPerceptionProvider(PerceptionProvider):
    def __init__(self, *, ocr_engine: PaddleOCREngine | None = None) -> None:
        self._ocr_engine = ocr_engine

    def build_scene(self, screenshot: bytes, tree: Any) -> Scene:
        width, height = infer_png_size(screenshot)
        try:
            ui_elements = parse_ui_tree(tree)
        except Exception:
            logger.warning(
                "UI tree parsing failed; continuing without UI tree", exc_info=True
            )
            ui_elements = []
        return build_fused_scene(
            screen_width=width,
            screen_height=height,
            ui_elements=ui_elements,
            ocr_elements=run_ocr(screenshot, engine=self._ocr_engine),
        )


class NullPerceptionProvider(PerceptionProvider):
    def build_scene(self, screenshot: bytes, tree: Any) -> Scene:
        width, height = infer_png_size(screenshot)
        return Scene(width=width, height=height, elements=[])