agentic-mobile-control/api/console.py

from __future__ import annotations

import base64
from pathlib import Path
from typing import Any
from uuid import uuid4

from device.manager import DeviceManager
from driver.registry import build_driver_factory
from pydantic import BaseModel, Field
from runtime.task import TaskRunner
from storage.device_config import DeviceConfigStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline


class RegisterDeviceRequest(BaseModel):
    driver_type: str
    connection_info: dict[str, Any] = Field(default_factory=dict)
    name: str | None = None


class RuntimeConfigRequest(BaseModel):
    max_steps: int


class ConsoleService:
    """API-local typed operations shared by JSON and HTML console routers.

    Stays in the ``api`` layer: it only depends on the existing injected
    Runtime state (device manager, metadata store, timeline, config store,
    task runner) and does not introduce HTTP/UI concerns below this layer.
    Page and JSON handlers receive the same instance so validation, driver
    allow-listing, persistence, and error mapping cannot drift apart.
    """

    def __init__(
        self,
        *,
        device_manager: DeviceManager,
        metadata_store: TaskMetadataStore,
        timeline: Timeline,
        config_store: DeviceConfigStore,
        task_runner: TaskRunner,
    ) -> None:
        self._device_manager = device_manager
        self._metadata_store = metadata_store
        self._timeline = timeline
        self._config_store = config_store
        self._task_runner = task_runner

    def list_devices(self) -> list[dict[str, Any]]:
        return [device.to_dict() for device in self._device_manager.list_devices()]

    def register_device(
        self,
        *,
        driver_type: str,
        connection_info: dict[str, Any],
        name: str | None,
    ) -> dict[str, Any]:
        driver_factory = build_driver_factory(driver_type, connection_info)
        device_id = uuid4().hex
        device = self._device_manager.register_device(
            device_id,
            driver_factory,
            name=name,
            driver_type=driver_type,
            connection_info=connection_info,
        )
        self._config_store.add(
            device_id=device_id,
            name=name,
            driver_type=driver_type,
            connection_info=connection_info,
        )
        return device.to_dict()

    def unregister_device(self, device_id: str) -> None:
        if not self._has_device(device_id):
            raise KeyError("device not found")
        self._device_manager.unregister_device(device_id)
        self._config_store.remove(device_id)

    def list_tasks(
        self,
        *,
        device_id: str | None = None,
        status: str | None = None,
    ) -> list[dict[str, Any]]:
        rows = self._metadata_store.list_tasks()
        if device_id is not None:
            rows = [row for row in rows if row["device_id"] == device_id]
        if status is not None:
            rows = [row for row in rows if row["status"] == status]
        return rows

    def get_task(self, task_id: str) -> dict[str, Any]:
        task = self._metadata_store.get_task(task_id)
        if task is None:
            raise KeyError("task not found")
        return task

    def get_task_timeline(self, task_id: str) -> list[dict[str, Any]]:
        if self._metadata_store.get_task(task_id) is None:
            raise KeyError("task not found")
        return [
            self._inline_screenshot(record) for record in self._timeline.read(task_id)
        ]

    def get_runtime_config(self) -> dict[str, int]:
        return {"max_steps": self._runner_max_steps()}

    def update_runtime_config(self, *, max_steps: int) -> dict[str, int]:
        if max_steps <= 0:
            raise ValueError("max_steps must be positive")
        self._set_runner_max_steps(max_steps)
        self._config_store.set_setting("max_steps", max_steps)
        return {"max_steps": max_steps}

    def _has_device(self, device_id: str) -> bool:
        return any(
            device.id == device_id for device in self._device_manager.list_devices()
        )

    @staticmethod
    def _inline_screenshot(record: dict[str, Any]) -> dict[str, Any]:
        payload = dict(record)
        screenshot_path = payload.get("screenshot_path")
        if screenshot_path:
            path = Path(str(screenshot_path))
            if path.exists():
                payload["image_base64"] = base64.b64encode(path.read_bytes()).decode(
                    "ascii"
                )
        return payload

    def _runner_max_steps(self) -> int:
        config = getattr(self._task_runner, "config", None)
        if config is None or not hasattr(config, "max_steps"):
            raise RuntimeError("task runner config unavailable")
        return int(config.max_steps)

    def _set_runner_max_steps(self, max_steps: int) -> None:
        config = getattr(self._task_runner, "config", None)
        if config is None or not hasattr(config, "max_steps"):
            raise RuntimeError("task runner config unavailable")
        config.max_steps = max_steps


def create_console_router(service: ConsoleService) -> Any:
    from fastapi import APIRouter, HTTPException, Response, status

    router = APIRouter(prefix="/console", tags=["console"])

    @router.get("/devices")
    def devices() -> list[dict[str, Any]]:
        return service.list_devices()

    @router.post("/devices", status_code=status.HTTP_201_CREATED)
    def register_device(request: RegisterDeviceRequest) -> dict[str, Any]:
        try:
            return service.register_device(
                driver_type=request.driver_type,
                connection_info=request.connection_info,
                name=request.name,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    @router.delete(
        "/devices/{device_id}",
        status_code=status.HTTP_204_NO_CONTENT,
        response_model=None,
    )
    def unregister_device(device_id: str) -> Response:
        try:
            service.unregister_device(device_id)
        except KeyError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        return Response(status_code=status.HTTP_204_NO_CONTENT)

    @router.get("/tasks")
    def tasks(
        device_id: str | None = None,
        status: str | None = None,
    ) -> list[dict[str, Any]]:
        return service.list_tasks(device_id=device_id, status=status)

    @router.get("/tasks/{task_id}")
    def task_detail(task_id: str) -> dict[str, Any]:
        try:
            return service.get_task(task_id)
        except KeyError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc

    @router.get("/tasks/{task_id}/timeline")
    def task_timeline(task_id: str) -> list[dict[str, Any]]:
        try:
            return service.get_task_timeline(task_id)
        except KeyError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc

    @router.get("/config")
    def runtime_config() -> dict[str, int]:
        return service.get_runtime_config()

    @router.put("/config")
    def update_runtime_config(
        request: RuntimeConfigRequest,
    ) -> dict[str, int]:
        try:
            return service.update_runtime_config(max_steps=request.max_steps)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    return router