agentic-mobile-control/api/console_web.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from api.console import ConsoleService
from driver.registry import SUPPORTED_DRIVER_TYPES
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from jinja2 import Environment, FileSystemLoader, select_autoescape

_TEMPLATE_DIR = Path(__file__).parent / "templates" / "runtime_console"
_STATIC_DIR = Path(__file__).parent / "static" / "runtime_console"

_ENV = Environment(
    loader=FileSystemLoader(str(_TEMPLATE_DIR)),
    autoescape=select_autoescape(["html", "xml"]),
    auto_reload=False,
)

_TASK_STATUSES = ["created", "running", "completed", "failed", "cancelled"]
_DEFAULT_FORM_VALUES: dict[str, str] = {
    "name": "",
    "driver_type": "",
    "server_url": "",
    "udid": "",
    "wda_local_port": "",
}


def _render(request: Request, template_name: str, **context: Any) -> str:
    return _ENV.get_template(template_name).render(
        url_for=request.url_for,
        **context,
    )


def _device_names(service: ConsoleService) -> dict[str, str]:
    return {
        device["id"]: device.get("name") or device["id"]
        for device in service.list_devices()
    }


def _default_driver_type() -> str:
    if not _DEFAULT_FORM_VALUES["driver_type"] and SUPPORTED_DRIVER_TYPES:
        _DEFAULT_FORM_VALUES["driver_type"] = next(iter(SUPPORTED_DRIVER_TYPES))
    return _DEFAULT_FORM_VALUES["driver_type"]


def _config_context(
    request: Request,
    service: ConsoleService,
    *,
    form_values: dict[str, str] | None = None,
    device_error: str = "",
    device_added: str = "",
    config_error: str = "",
    config_saved: bool = False,
    max_steps_override: int | None = None,
) -> str:
    form = dict(_DEFAULT_FORM_VALUES)
    form["driver_type"] = _default_driver_type()
    if form_values:
        form.update(form_values)
    config = service.get_runtime_config()
    return _render(
        request,
        "config.html",
        active_section="config",
        devices=service.list_devices(),
        max_steps=max_steps_override
        if max_steps_override is not None
        else config["max_steps"],
        supported_driver_types=list(SUPPORTED_DRIVER_TYPES),
        form_values=form,
        device_error=device_error,
        device_added=device_added,
        config_error=config_error,
        config_saved=config_saved,
    )


def create_console_web_router(service: ConsoleService) -> APIRouter:
    router = APIRouter(prefix="/ui", tags=["console-ui"])

    @router.get("/", name="runtime_console_dashboard", include_in_schema=False)
    def dashboard(request: Request) -> HTMLResponse:
        devices = service.list_devices()
        tasks = service.list_tasks()
        running = sum(1 for task in tasks if task.get("status") == "running")
        failed = sum(1 for task in tasks if task.get("status") == "failed")
        return HTMLResponse(
            _render(
                request,
                "dashboard.html",
                active_section="dashboard",
                devices=devices,
                device_count=len(devices),
                task_count=len(tasks),
                running_tasks=running,
                failed_tasks=failed,
            )
        )

    @router.get(
        "/_status_fragment",
        name="runtime_console_status_fragment",
        include_in_schema=False,
    )
    def status_fragment(request: Request) -> HTMLResponse:
        devices = service.list_devices()
        tasks = service.list_tasks()
        running = sum(1 for task in tasks if task.get("status") == "running")
        failed = sum(1 for task in tasks if task.get("status") == "failed")
        return HTMLResponse(
            _render(
                request,
                "_status_fragment.html",
                devices=devices,
                device_count=len(devices),
                running_tasks=running,
                failed_tasks=failed,
            )
        )

    @router.get("/tasks", name="runtime_console_tasks", include_in_schema=False)
    def tasks(
        request: Request,
        device_id: str | None = None,
        status: str | None = None,
    ) -> HTMLResponse:
        service_tasks = service.list_tasks(device_id=device_id, status=status)
        return HTMLResponse(
            _render(
                request,
                "tasks.html",
                active_section="tasks",
                tasks=service_tasks,
                devices=service.list_devices(),
                device_names=_device_names(service),
                task_statuses=_TASK_STATUSES,
                selected_device_id=device_id or "",
                selected_status=status or "",
            )
        )

    @router.get(
        "/tasks/{task_id}",
        name="runtime_console_task_detail",
        include_in_schema=False,
    )
    def task_detail(
        request: Request,
        task_id: str,
        step: int | None = None,
    ) -> HTMLResponse:
        try:
            task = service.get_task(task_id)
            timeline = service.get_task_timeline(task_id)
        except KeyError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        if timeline:
            index = 0
            if step is not None:
                index = max(0, min(step, len(timeline) - 1))
            current_step = timeline[index]
        else:
            index = 0
            current_step = None
        device_name = _device_names(service).get(
            task.get("device_id", ""),
            task.get("device_id", ""),
        )
        return HTMLResponse(
            _render(
                request,
                "task_detail.html",
                active_section="tasks",
                task=task,
                device_name=device_name,
                timeline=timeline,
                current_step=current_step if current_step is not None else {},
                current_step_index=index,
            )
        )

    @router.get("/config", name="runtime_console_config", include_in_schema=False)
    def config(request: Request) -> HTMLResponse:
        return HTMLResponse(_config_context(request, service))

    @router.post(
        "/config/devices",
        name="runtime_console_register_device",
        include_in_schema=False,
        response_model=None,
    )
    async def register_device(request: Request) -> HTMLResponse | RedirectResponse:
        form = await request.form()
        name = (form.get("name") or "").strip()
        driver_type = (form.get("driver_type") or "").strip()
        server_url = (form.get("server_url") or "").strip()
        udid = (form.get("udid") or "").strip()
        wda_local_port_raw = (form.get("wda_local_port") or "").strip()

        connection_info: dict[str, Any] = {}
        if server_url:
            connection_info["server_url"] = server_url
        if udid:
            connection_info["udid"] = udid
        if wda_local_port_raw:
            try:
                connection_info["wda_local_port"] = int(wda_local_port_raw)
            except ValueError:
                return HTMLResponse(
                    _config_context(
                        request,
                        service,
                        form_values={
                            "name": name,
                            "driver_type": driver_type,
                            "server_url": server_url,
                            "udid": udid,
                            "wda_local_port": wda_local_port_raw,
                        },
                        device_error="wda_local_port must be a number",
                    ),
                    status_code=400,
                )

        try:
            service.register_device(
                driver_type=driver_type,
                connection_info=connection_info,
                name=name or None,
            )
        except ValueError as exc:
            return HTMLResponse(
                _config_context(
                    request,
                    service,
                    form_values={
                        "name": name,
                        "driver_type": driver_type,
                        "server_url": server_url,
                        "udid": udid,
                        "wda_local_port": wda_local_port_raw,
                    },
                    device_error=str(exc),
                ),
                status_code=400,
            )
        return RedirectResponse(
            url=request.url_for("runtime_console_config"),
            status_code=303,
        )

    @router.post(
        "/config/devices/{device_id}/delete",
        name="runtime_console_remove_device",
        include_in_schema=False,
        response_model=None,
    )
    async def remove_device(
        request: Request,
        device_id: str,
    ) -> HTMLResponse | RedirectResponse:
        try:
            service.unregister_device(device_id)
        except KeyError as exc:
            return HTMLResponse(
                _config_context(request, service, device_error=str(exc)),
                status_code=404,
            )
        return RedirectResponse(
            url=request.url_for("runtime_console_config"),
            status_code=303,
        )

    @router.post(
        "/config/max-steps",
        name="runtime_console_update_max_steps",
        include_in_schema=False,
        response_model=None,
    )
    async def update_max_steps(request: Request) -> HTMLResponse | RedirectResponse:
        form = await request.form()
        raw = (form.get("max_steps") or "").strip()
        try:
            max_steps = int(raw)
        except ValueError:
            return HTMLResponse(
                _config_context(
                    request,
                    service,
                    config_error="max_steps must be an integer",
                    max_steps_override=raw if raw else None,
                ),
                status_code=400,
            )
        try:
            service.update_runtime_config(max_steps=max_steps)
        except ValueError as exc:
            return HTMLResponse(
                _config_context(
                    request,
                    service,
                    config_error=str(exc),
                    max_steps_override=max_steps,
                ),
                status_code=400,
            )
        return RedirectResponse(
            url=request.url_for("runtime_console_config"),
            status_code=303,
        )

    return router


def mount_console_assets(app: Any) -> None:
    """Mount the package-owned static assets under ``/ui/assets``.

    Kept as a separate helper so ``api.rest`` can mount assets with a stable
    route name that templates reference via ``url_for``.
    """
    app.mount(
        "/ui/assets",
        StaticFiles(directory=str(_STATIC_DIR)),
        name="runtime_console_assets",
    )