app.py

"""
FastAPI server for the OpenX Rerun Viewer HF Space.
Runs the selection UI and serves .rrd files.
When running locally, also starts a local Rerun web viewer to avoid mixed-content issues.
"""
import json
import os
import subprocess
import sys
import time
from pathlib import Path

from fastapi import FastAPI
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI(title="OpenX Rerun Viewer")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "HEAD", "OPTIONS"],
    allow_headers=["*"],
    expose_headers=[
        "Content-Length", "Content-Type", "Accept-Ranges",
        "Cross-Origin-Resource-Policy",
    ],
)

DATA_DIR = Path(__file__).parent / "data"
MANIFEST_PATH = DATA_DIR / "manifest.json"

IS_LOCAL = "SPACE_ID" not in os.environ and "HF_SPACE" not in os.environ


def load_manifest():
    if MANIFEST_PATH.exists():
        return json.loads(MANIFEST_PATH.read_text())
    return []


@app.get("/", response_class=HTMLResponse)
async def index():
    html = (Path(__file__).parent / "templates" / "index.html").read_text()
    # Inject the viewer base URL: use local viewer for dev, app.rerun.io for production
    if IS_LOCAL:
        viewer_base = "http://localhost:9090"
    else:
        viewer_base = "https://app.rerun.io/version/0.26.2"
    html = html.replace("{{VIEWER_BASE}}", viewer_base)
    return HTMLResponse(html)


@app.get("/api/datasets")
async def list_datasets():
    manifest = load_manifest()
    return [
        {
            "slug": m["slug"],
            "robot_type": m["robot_type"],
            "fps": m.get("fps", 0),
            "camera_keys": m.get("camera_keys", []),
            "camera_labels": m.get("camera_labels", []),
            "num_episodes_total": m.get("num_episodes_total", 0),
            "state_dim": m.get("state_dim", 0),
            "action_dim": m.get("action_dim", 0),
            "episodes_available": [
                {"index": ep["index"], "frames": ep.get("frames", "?")}
                for ep in m.get("episodes_available", [])
            ],
        }
        for m in manifest
    ]


@app.get("/api/datasets/{slug}/episodes")
async def list_episodes(slug: str):
    manifest = load_manifest()
    for m in manifest:
        if m["slug"] == slug:
            return m["episodes_available"]
    return []


@app.api_route("/data/{slug}/episode_{index:str}.rrd", methods=["GET", "HEAD"])
async def serve_rrd(slug: str, index: str):
    filepath = DATA_DIR / slug / f"episode_{index}.rrd"
    if not filepath.exists():
        return HTMLResponse(status_code=404, content="Episode not found")
    return FileResponse(
        filepath,
        media_type="application/octet-stream",
        headers={
            "Accept-Ranges": "bytes",
            "Access-Control-Allow-Origin": "*",
            "Cross-Origin-Resource-Policy": "cross-origin",
            "Cache-Control": "public, max-age=3600",
        },
    )


static_dir = Path(__file__).parent / "static"
if static_dir.exists():
    app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")


def start_local_viewer():
    """Start Rerun web viewer on port 9090 for local development."""
    import rerun as rr
    rr.serve_web_viewer(web_port=9090, open_browser=False)
    # serve_web_viewer returns immediately after starting the server in a thread


if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("PORT", 8000))

    if IS_LOCAL:
        print("Local mode: starting Rerun web viewer on http://localhost:9090")
        start_local_viewer()
        time.sleep(0.5)

    uvicorn.run(app, host="0.0.0.0", port=port)