from __future__ import annotations

from pathlib import Path
from typing import Any

from .runtime.model import LoadedUIActionPolicy
from .schema import UIAction, UIElement


import os
DEFAULT_CHECKPOINT = Path(os.environ.get("LEO_UI7M_CHECKPOINT", Path(__file__).resolve().parents[1] / "checkpoints" / "leo_ui7m_v4_best.pt"))


class UIActionModel:
    """LEO-UI7M local inference API."""

    def __init__(self, checkpoint_path: str | Path | None = None, device: str = "cpu") -> None:
        self.checkpoint_path = Path(checkpoint_path) if checkpoint_path else DEFAULT_CHECKPOINT
        self.device = device
        self.policy = LoadedUIActionPolicy(self.checkpoint_path, device=device)

    def predict_action(
        self,
        goal: str,
        elements: list[UIElement | dict[str, Any]],
        history: list[dict[str, Any]] | None = None,
        step_index: int = 0,
    ) -> UIAction:
        normalized = [
            e if isinstance(e, dict) else e.__dict__
            for e in elements
        ]
        raw = self.policy.predict_raw(
            goal=goal,
            elements=normalized,
            history=history or [],
            step_index=step_index,
        )
        return UIAction(
            action=raw["action"],
            target_element_id=raw.get("target_element_id"),
            confidence=raw.get("confidence"),
            raw=raw,
        )