agentic-mobile-control/core/models.py

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any, Literal
from uuid import uuid4

DeviceStatus = Literal["idle", "busy", "offline", "error"]
TaskStatus = Literal["created", "running", "completed", "failed", "cancelled"]
StepStatus = Literal["pending", "running", "completed", "failed"]


def utc_now() -> datetime:
    return datetime.now(UTC)


@dataclass(frozen=True)
class Bounds:
    x: float
    y: float
    width: float
    height: float

    @property
    def right(self) -> float:
        return self.x + self.width

    @property
    def bottom(self) -> float:
        return self.y + self.height

    @property
    def center(self) -> tuple[float, float]:
        return (self.x + self.width / 2, self.y + self.height / 2)

    def to_dict(self) -> dict[str, float]:
        return {
            "x": self.x,
            "y": self.y,
            "width": self.width,
            "height": self.height,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Bounds":
        return cls(
            x=float(data["x"]),
            y=float(data["y"]),
            width=float(data["width"]),
            height=float(data["height"]),
        )


@dataclass
class Device:
    id: str
    status: DeviceStatus = "idle"
    name: str | None = None
    driver_type: str = "wda"
    connection_info: dict[str, Any] = field(default_factory=dict)
    capability_tags: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "status": self.status,
            "driver_type": self.driver_type,
            "connection_info": dict(self.connection_info),
            "capability_tags": list(self.capability_tags),
        }


_STATE_FIELDS = ("enabled", "clickable", "selected", "checked", "focused")
_COLOR_FIELDS = ("foreground_color", "background_color")


@dataclass
class SceneElement:
    id: str
    type: str
    bounds: Bounds
    text: str | None = None
    confidence: float | None = None
    source: str | None = None
    # Accessibility-tree interaction state, when the platform reports it.
    # None means "not reported by this platform/element", not "false".
    enabled: bool | None = None
    clickable: bool | None = None
    selected: bool | None = None
    checked: bool | None = None
    focused: bool | None = None
    # Text/background color sampled from the screenshot pixels under the OCR
    # box ("#rrggbb"). None means unavailable (not OCR-sourced, or sampling
    # failed), not "no color".
    foreground_color: str | None = None
    background_color: str | None = None

    @property
    def center(self) -> tuple[float, float]:
        return self.bounds.center

    def to_dict(self) -> dict[str, Any]:
        data: dict[str, Any] = {
            "id": self.id,
            "type": self.type,
            "text": self.text,
            "bounds": self.bounds.to_dict(),
            "confidence": self.confidence,
        }
        if self.source:
            data["source"] = self.source
        for field_name in _STATE_FIELDS + _COLOR_FIELDS:
            value = getattr(self, field_name)
            if value is not None:
                data[field_name] = value
        return data

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SceneElement":
        return cls(
            id=str(data["id"]),
            type=str(data.get("type") or "unknown"),
            text=data.get("text"),
            bounds=Bounds.from_dict(data["bounds"]),
            confidence=data.get("confidence"),
            source=data.get("source"),
            enabled=data.get("enabled"),
            clickable=data.get("clickable"),
            selected=data.get("selected"),
            checked=data.get("checked"),
            focused=data.get("focused"),
            foreground_color=data.get("foreground_color"),
            background_color=data.get("background_color"),
        )


@dataclass
class Scene:
    width: int
    height: int
    elements: list[SceneElement] = field(default_factory=list)
    # Keep raw OCR observations for local execution evidence without duplicating
    # them in the normalized, LLM-facing scene payload.
    ocr_elements: list[SceneElement] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "screen": {"width": self.width, "height": self.height},
            "elements": [element.to_dict() for element in self.elements],
        }

    def ocr_results_to_dict(self) -> list[dict[str, Any]]:
        return [element.to_dict() for element in self.ocr_elements]

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Scene":
        screen = data.get("screen") or {}
        return cls(
            width=int(screen.get("width") or data.get("width") or 0),
            height=int(screen.get("height") or data.get("height") or 0),
            elements=[
                SceneElement.from_dict(element) for element in data.get("elements", [])
            ],
            ocr_elements=[
                SceneElement.from_dict(element)
                for element in data.get("ocr_elements", [])
            ],
        )


@dataclass
class Task:
    goal: str
    device_id: str
    id: str = field(default_factory=lambda: uuid4().hex)
    status: TaskStatus = "created"
    created_at: datetime = field(default_factory=utc_now)
    updated_at: datetime = field(default_factory=utc_now)
    completed_at: datetime | None = None
    failure_reason: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "goal": self.goal,
            "device_id": self.device_id,
            "status": self.status,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "completed_at": self.completed_at.isoformat()
            if self.completed_at
            else None,
            "failure_reason": self.failure_reason,
        }


@dataclass
class Step:
    description: str
    action: str
    args: dict[str, Any] = field(default_factory=dict)
    id: str = field(default_factory=lambda: uuid4().hex)
    status: StepStatus = "pending"
    result: dict[str, Any] | None = None
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "description": self.description,
            "action": self.action,
            "args": dict(self.args),
            "status": self.status,
            "result": self.result,
            "error": self.error,
        }