app.py

"""Grabette fleet — operator dashboard + command broker (Docker HF Space).

A free-tier-friendly Docker Space. It is self-contained (no `grabette` import) so
it deploys standalone to HF. Responsibilities:

  * Operator UI + login via HF Spaces native OAuth (`hf_oauth: true`).
  * Command broker: an in-memory, per-owner device registry + command queue.
  * Device auth: devices call with `Authorization: Bearer <hf_token>`; we resolve
    the owner via `whoami` (cached) and group devices by HF identity.

Transport is short-polling (devices GET /api/devices/poll every couple seconds).
That polling traffic doubles as the keep-alive that stops a free Space sleeping
(sleep is timed from the last request). State is in-memory, so a restart drops
it — devices simply re-register on their next poll; durable data lives in the
device's own HF datasets, not here.

Designed to be duplicated per user: one Space = one owner = one fleet.
"""

from __future__ import annotations

import os
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Optional

from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.responses import HTMLResponse
from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth, whoami
from pydantic import BaseModel

ONLINE_WINDOW = 15.0  # device is "online" if seen within this many seconds
WHOAMI_TTL = 300.0


# --- device identity (Bearer token -> HF username, cached) -------------------
_whoami_cache: dict[str, tuple[str, float]] = {}


def verify_user(token: str) -> str:
    now = time.time()
    hit = _whoami_cache.get(token)
    if hit and hit[1] > now:
        return hit[0]
    try:
        info = whoami(token=token)
    except Exception as e:  # noqa: BLE001
        raise ValueError(f"Invalid HF token: {e}") from e
    name = (info or {}).get("name", "")
    if not name:
        raise ValueError("Could not resolve HF username")
    _whoami_cache[token] = (name, now + WHOAMI_TTL)
    return name


async def device_owner(authorization: Optional[str] = Header(None)) -> str:
    """FastAPI dep: resolve the device's owner from its Bearer HF token."""
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(401, "Missing 'Authorization: Bearer <hf_token>'")
    try:
        return verify_user(authorization[len("Bearer ") :].strip())
    except ValueError as e:
        raise HTTPException(401, str(e)) from e


def operator_name(request: Request) -> str:
    """Resolve the logged-in operator's HF username, or 401."""
    info = parse_huggingface_oauth(request)
    if info is None:
        raise HTTPException(401, "Not logged in")
    return info.user_info.preferred_username or info.user_info.name


# --- in-memory fleet state ---------------------------------------------------
@dataclass
class Command:
    id: str
    type: str
    args: dict[str, Any]
    status: str = "pending"  # pending | sent | done
    result: Optional[dict[str, Any]] = None
    created_at: float = field(default_factory=time.time)


@dataclass
class Device:
    device_id: str
    name: str
    capabilities: list[str]
    last_seen: float = field(default_factory=time.time)
    queue: list[Command] = field(default_factory=list)
    history: list[Command] = field(default_factory=list)

    @property
    def online(self) -> bool:
        return (time.time() - self.last_seen) < ONLINE_WINDOW


FLEET: dict[str, dict[str, Device]] = {}  # owner -> device_id -> Device


def _fleet_of(owner: str) -> dict[str, Device]:
    return FLEET.setdefault(owner, {})


# --- app + OAuth -------------------------------------------------------------
app = FastAPI(title="Grabette fleet")
attach_huggingface_oauth(app)  # adds /oauth/huggingface/{login,logout,callback}


# === device-facing API (Bearer auth) ========================================
class RegisterReq(BaseModel):
    device_id: str
    name: str = ""
    capabilities: list[str] = []


class ResultReq(BaseModel):
    device_id: str
    command_id: str
    result: dict[str, Any] = {}


@app.post("/api/devices/register")
async def register(req: RegisterReq, owner: str = Depends(device_owner)) -> dict[str, Any]:
    fleet = _fleet_of(owner)
    dev = fleet.get(req.device_id)
    if dev is None:
        dev = Device(req.device_id, req.name or req.device_id, req.capabilities)
        fleet[req.device_id] = dev
    else:
        dev.name = req.name or dev.name
        dev.capabilities = req.capabilities or dev.capabilities
    dev.last_seen = time.time()
    return {"status": "ok", "pending": len(dev.queue)}


@app.get("/api/devices/poll")
async def poll(device_id: str, owner: str = Depends(device_owner)) -> dict[str, Any]:
    dev = _fleet_of(owner).get(device_id)
    if dev is None:
        raise HTTPException(404, "Device not registered")
    dev.last_seen = time.time()
    pending = [c for c in dev.queue if c.status == "pending"]
    for c in pending:
        c.status = "sent"
    return {"commands": [{"id": c.id, "type": c.type, "args": c.args} for c in pending]}


@app.post("/api/devices/result")
async def result(req: ResultReq, owner: str = Depends(device_owner)) -> dict[str, str]:
    dev = _fleet_of(owner).get(req.device_id)
    if dev is None:
        raise HTTPException(404, "Device not registered")
    for c in dev.queue:
        if c.id == req.command_id:
            c.status = "done"
            c.result = req.result
            dev.queue.remove(c)
            dev.history.insert(0, c)
            del dev.history[20:]
            break
    return {"status": "ok"}


# === operator-facing API (session OAuth) =====================================
class DispatchReq(BaseModel):
    device_id: str
    type: str
    args: dict[str, Any] = {}


@app.get("/api/fleet/me")
async def me(request: Request) -> dict[str, Any]:
    info = parse_huggingface_oauth(request)
    if info is None:
        return {"logged_in": False}
    return {"logged_in": True, "username": info.user_info.preferred_username or info.user_info.name}


@app.get("/api/fleet/devices")
async def list_devices(request: Request) -> dict[str, Any]:
    owner = operator_name(request)
    devices = _fleet_of(owner).values()
    return {
        "owner": owner,
        "devices": [
            {
                "device_id": d.device_id,
                "name": d.name,
                "online": d.online,
                "capabilities": d.capabilities,
                "pending": len([c for c in d.queue if c.status != "done"]),
                "history": [
                    {"type": c.type, "args": c.args, "result": c.result} for c in d.history[:5]
                ],
            }
            for d in devices
        ],
    }


@app.post("/api/fleet/dispatch")
async def dispatch(req: DispatchReq, request: Request) -> dict[str, Any]:
    owner = operator_name(request)
    dev = _fleet_of(owner).get(req.device_id)
    if dev is None:
        raise HTTPException(404, "Device not found in your fleet")
    cmd = Command(id=uuid.uuid4().hex[:12], type=req.type, args=req.args)
    dev.queue.append(cmd)
    return {"status": "queued", "command_id": cmd.id}


# === UI ======================================================================
@app.get("/")
async def index() -> HTMLResponse:
    return HTMLResponse(_INDEX)


_INDEX = """<!DOCTYPE html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Grabette fleet</title><style>
 body{font-family:-apple-system,system-ui,sans-serif;background:linear-gradient(135deg,#1a1a2e,#16213e);
  color:#fff;min-height:100vh;margin:0;padding:2rem;display:flex;justify-content:center}
 .wrap{width:560px}h1{font-size:1.3rem}h2{font-size:1rem;margin:0 0 .8rem}
 .card{background:rgba(255,255,255,.06);padding:1.2rem;border-radius:14px;margin-bottom:1rem}
 a.btn,button{padding:.55rem 1rem;border:0;border-radius:8px;cursor:pointer;font-weight:600;text-decoration:none;display:inline-block}
 .primary{background:#ffcc4d;color:#1a1a2e}.logout{background:#ef4444;color:#fff}
 .muted{color:#a0aec0;font-size:.82rem}
 table{width:100%;border-collapse:collapse;font-size:.85rem}td,th{text-align:left;padding:.35rem;border-bottom:1px solid #334}
 .dot{display:inline-block;width:.6rem;height:.6rem;border-radius:50%}.on{background:#10b981}.off{background:#ef4444}
 fieldset:disabled{opacity:.45}fieldset{border:0;padding:0;margin:0}code{color:#ffcc4d}
</style></head><body><div class="wrap">
<h1>🤖 Grabette fleet <span class="muted">— operator dashboard</span></h1>
<div class="card"><h2>HuggingFace login</h2><div id="who">Checking…</div></div>
<fieldset class="card" id="gated" disabled>
 <h2>Your devices</h2>
 <table id="devs"><thead><tr><th></th><th>device</th><th>commands</th><th>last result</th></tr></thead><tbody></tbody></table>
 <div class="muted" id="empty" style="margin-top:.6rem">No devices reporting yet. Start one with <code>device_app.py</code>.</div>
</fieldset></div><script>
const $=id=>document.getElementById(id);let loggedIn=false;
function btn(id,type,args){return `<button class="primary" onclick='dispatch("${id}","${type}",${JSON.stringify(args||{})})'>${type}</button>`;}
async function dispatch(device_id,type,args){
 await fetch('/api/fleet/dispatch',{method:'POST',headers:{'Content-Type':'application/json'},
  body:JSON.stringify({device_id,type,args})});refresh();}
async function checkLogin(){
 const m=await(await fetch('/api/fleet/me')).json();loggedIn=m.logged_in;
 if(loggedIn){$('who').innerHTML=`Signed in as <b>${m.username}</b> <a class="btn logout" href="/oauth/huggingface/logout">Logout</a>`;}
 else{$('who').innerHTML=`<a class="btn primary" href="/oauth/huggingface/login">Sign in with HuggingFace</a>`;}
 $('gated').disabled=!loggedIn;}
async function refresh(){
 if(!loggedIn)return;
 const r=await fetch('/api/fleet/devices');if(!r.ok)return;const j=await r.json();
 const tb=$('devs').querySelector('tbody');tb.innerHTML='';
 $('empty').style.display=j.devices.length?'none':'block';
 for(const d of j.devices){
  const last=d.history[0];const lr=last?`${last.type} → ${JSON.stringify(last.result)}`:'—';
  const caps=(d.capabilities&&d.capabilities.length)?d.capabilities:['power_on','power_off','move'];
  const acts=caps.map(c=>c==='move'
    ?`<button class="primary" onclick='dispatch("${d.device_id}","move",{position:Math.floor(Math.random()*100)})'>move</button>`
    :btn(d.device_id,c)).join(' ');
  tb.insertAdjacentHTML('beforeend',
   `<tr><td><span class="dot ${d.online?'on':'off'}"></span></td>
    <td><b>${d.name}</b><br><span class="muted">${d.device_id}${d.pending?` · ${d.pending} pending`:''}</span></td>
    <td>${acts}</td><td class="muted">${lr}</td></tr>`);}}
async function tick(){await checkLogin();await refresh();}
tick();setInterval(tick,3000);
</script></body></html>"""