foodlinkk-command-center/browser-agent/app/pa_live.py

"""PA multi-site live status — 4 browser slots for Cockpit."""
from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

PA_SITES = ("Airbnb", "Booking.com", "DuckDuckGo", "HolidayCheck")

_state: dict[str, Any] = {
    "job_id": None,
    "query": "",
    "chat_id": None,
    "user_name": "",
    "status": "idle",
    "updated_at": None,
    "slots": {
        site: {
            "label": site,
            "status": "idle",
            "url": None,
            "title": None,
            "error": None,
            "index": i + 1,
            "total": len(PA_SITES),
            "screenshot_b64": None,
            "updated_at": None,
        }
        for i, site in enumerate(PA_SITES)
    },
}


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _slot(label: str) -> dict[str, Any]:
    slots = _state["slots"]
    if label not in slots:
        slots[label] = {
            "label": label,
            "status": "idle",
            "url": None,
            "title": None,
            "error": None,
            "index": 0,
            "total": len(PA_SITES),
            "screenshot_b64": None,
            "updated_at": _now(),
        }
    return slots[label]


def start_job(
    job_id: str,
    query: str,
    *,
    chat_id: int | None = None,
    user_name: str = "",
    sites: list[str] | None = None,
) -> None:
    site_list = sites or list(PA_SITES)
    _state["job_id"] = job_id
    _state["query"] = query[:200]
    _state["chat_id"] = chat_id
    _state["user_name"] = user_name
    _state["status"] = "running"
    _state["updated_at"] = _now()
    _state["slots"] = {
        site: {
            "label": site,
            "status": "waiting",
            "url": None,
            "title": None,
            "error": None,
            "index": i + 1,
            "total": len(site_list),
            "screenshot_b64": None,
            "updated_at": _now(),
        }
        for i, site in enumerate(site_list)
    }


def set_comparing() -> None:
    _state["status"] = "comparing"
    _state["updated_at"] = _now()


def finish_job() -> None:
    _state["status"] = "done"
    _state["updated_at"] = _now()


def fail_job(error: str) -> None:
    _state["status"] = "failed"
    _state["updated_at"] = _now()
    _state["error"] = error[:200]


def update_slot(
    label: str,
    *,
    status: str,
    url: str | None = None,
    title: str | None = None,
    error: str | None = None,
    index: int | None = None,
    total: int | None = None,
    screenshot_b64: str | None = None,
) -> None:
    slot = _slot(label)
    slot["status"] = status
    slot["updated_at"] = _now()
    if url is not None:
        slot["url"] = url
    if title is not None:
        slot["title"] = title
    if error is not None:
        slot["error"] = error[:300]
    if index is not None:
        slot["index"] = index
    if total is not None:
        slot["total"] = total
    if screenshot_b64 is not None:
        slot["screenshot_b64"] = screenshot_b64
    _state["updated_at"] = _now()


def get_live() -> dict[str, Any]:
    out = {
        "job_id": _state.get("job_id"),
        "query": _state.get("query"),
        "chat_id": _state.get("chat_id"),
        "user_name": _state.get("user_name"),
        "status": _state.get("status", "idle"),
        "updated_at": _state.get("updated_at"),
        "sites": list(PA_SITES),
        "slots": [],
    }
    for site in PA_SITES:
        slot = dict(_state["slots"].get(site) or _slot(site))
        slot.pop("screenshot_b64", None)
        slot["has_screenshot"] = bool(
            (_state["slots"].get(site) or {}).get("screenshot_b64")
        )
        out["slots"].append(slot)
    return out


def get_slot_screenshot(label: str) -> bytes | None:
    slot = _state["slots"].get(label)
    if not slot or not slot.get("screenshot_b64"):
        return None
    import base64

    return base64.b64decode(slot["screenshot_b64"])