app.py

#!/usr/bin/env python3
"""Public review service for the WorldModelBench Hugging Face Space."""

from __future__ import annotations

import io
import json
import os
import threading
import uuid
from contextlib import closing
from datetime import datetime, timezone
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import parse_qs, urlparse
from urllib.request import Request, urlopen

from huggingface_hub import CommitOperationAdd, HfApi, get_token


ROOT = Path(__file__).resolve().parent
STATIC_ROOT = ROOT / "static"
MEDIA_CACHE_ROOT = ROOT / "media-cache"
CATALOG_PATH = STATIC_ROOT / "data" / "worldmodelbench-catalog.js"
DATASET_REPO = os.environ.get("WORLDMODELBENCH_REPO", "Xinyun-Liu/WorldModelBench")
HF_TOKEN = os.environ.get("HF_TOKEN") or get_token() or ""
MAX_BODY_BYTES = 32_000
MIN_SUBMISSION_INTERVAL = 0.8
last_submission_by_ip: dict[str, float] = {}
rate_lock = threading.Lock()


def read_catalog() -> dict[str, dict[str, str]]:
    raw = CATALOG_PATH.read_text(encoding="utf-8")
    payload = raw.split("=", 1)[1].strip().rstrip(";")
    rows = json.loads(payload)
    return {row["videoPath"]: row for row in rows if isinstance(row, dict) and row.get("videoPath")}


CATALOG = read_catalog()


def cached_media_path(video_path: str) -> Path | None:
    """Return a cache entry only when it stays beneath the cache root."""
    cache_root = MEDIA_CACHE_ROOT.resolve()
    try:
        candidate = (cache_root / video_path).resolve()
        candidate.relative_to(cache_root)
    except ValueError:
        return None
    return candidate if candidate.is_file() else None


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def clean_string(value: object, field: str, maximum: int, required: bool = False) -> str:
    if not isinstance(value, str):
        raise ValueError(f"{field}_must_be_text")
    value = value.strip()
    if required and not value:
        raise ValueError(f"{field}_is_required")
    if len(value) > maximum:
        raise ValueError(f"{field}_is_too_long")
    return value


def validate_submission(payload: object) -> dict[str, str]:
    if not isinstance(payload, dict):
        raise ValueError("request_must_be_a_json_object")
    annotator = clean_string(payload.get("annotator"), "annotator", 80, required=True)
    video_path = clean_string(payload.get("video_path"), "video_path", 500, required=True)
    row = CATALOG.get(video_path)
    if row is None:
        raise ValueError("video_path_is_not_in_the_catalog")
    decision = clean_string(payload.get("decision", "pending"), "decision", 20, required=True)
    if decision not in {"pending", "keep", "reject"}:
        raise ValueError("invalid_decision")
    return {
        "annotator": annotator,
        "video_path": video_path,
        "sample_id": clean_string(row.get("id", ""), "sample_id", 300, required=True),
        "data_source": clean_string(row.get("sourceDataset", ""), "data_source", 200, required=True),
        "decision": decision,
        "track": clean_string(payload.get("track", ""), "track", 200),
        "scene": clean_string(payload.get("scene", ""), "scene", 300),
        "spatial_ability": clean_string(payload.get("spatial_ability", ""), "spatial_ability", 500),
        "perspective": clean_string(payload.get("perspective", ""), "perspective", 300),
        "text_caption": clean_string(payload.get("text_caption", ""), "text_caption", 4000),
    }


class ReviewHandler(SimpleHTTPRequestHandler):
    server_version = "WorldModelBenchReview/1.0"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=str(STATIC_ROOT), **kwargs)

    def end_headers(self) -> None:
        if urlparse(self.path).path == "/review.html":
            self.send_header("Cache-Control", "no-store")
        super().end_headers()

    def send_json(self, status: int, value: dict[str, object]) -> None:
        body = json.dumps(value, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def proxy_asset(self, url: str, send_body: bool = True) -> None:
        """Relay an HF resolve URL as a same-origin, range-capable response."""
        headers = {"User-Agent": "WorldModelBench-Review/1.0"}
        requested_range = self.headers.get("Range")
        if requested_range:
            headers["Range"] = requested_range
        try:
            with closing(urlopen(Request(url, headers=headers), timeout=90)) as upstream:
                self.send_response(upstream.status)
                for name in ("Content-Type", "Content-Length", "Content-Range", "Accept-Ranges", "Last-Modified", "ETag"):
                    value = upstream.headers.get(name)
                    if value:
                        self.send_header(name, value)
                self.send_header("Cache-Control", "public, max-age=300")
                self.end_headers()
                if send_body:
                    while chunk := upstream.read(1024 * 1024):
                        self.wfile.write(chunk)
        except HTTPError as error:
            self.send_json(error.code, {"error": "upstream_asset_unavailable"})
        except (OSError, URLError, TimeoutError):
            self.send_json(502, {"error": "upstream_asset_connection_failed"})

    def serve_cached_media(self, asset: Path, send_body: bool = True) -> None:
        """Serve a locally cached original MP4 with browser Range support."""
        total_size = asset.stat().st_size
        start, end, partial = 0, total_size - 1, False
        requested_range = self.headers.get("Range", "")
        if requested_range:
            try:
                unit, spec = requested_range.split("=", 1)
                if unit.strip() != "bytes" or "," in spec:
                    raise ValueError
                first, last = spec.strip().split("-", 1)
                if first:
                    start = int(first)
                    end = int(last) if last else total_size - 1
                elif last:
                    suffix = int(last)
                    if suffix <= 0:
                        raise ValueError
                    start = max(total_size - suffix, 0)
                else:
                    raise ValueError
                if start < 0 or start >= total_size or end < start:
                    raise ValueError
                end = min(end, total_size - 1)
                partial = True
            except (ValueError, TypeError):
                self.send_response(416)
                self.send_header("Content-Range", f"bytes */{total_size}")
                self.send_header("Content-Length", "0")
                self.end_headers()
                return

        length = end - start + 1
        self.send_response(206 if partial else 200)
        self.send_header("Content-Type", "video/mp4")
        self.send_header("Content-Length", str(length))
        self.send_header("Accept-Ranges", "bytes")
        if partial:
            self.send_header("Content-Range", f"bytes {start}-{end}/{total_size}")
        self.send_header("Cache-Control", "public, max-age=300")
        self.end_headers()
        if not send_body:
            return
        try:
            with asset.open("rb") as stream:
                stream.seek(start)
                remaining = length
                while remaining:
                    chunk = stream.read(min(1024 * 1024, remaining))
                    if not chunk:
                        break
                    self.wfile.write(chunk)
                    remaining -= len(chunk)
        except BrokenPipeError:
            return

    def proxy_catalog_asset(self, kind: str, query: dict[str, list[str]], send_body: bool = True) -> None:
        paths = query.get("path", [])
        if len(paths) != 1:
            self.send_json(400, {"error": "exactly_one_video_path_is_required"})
            return
        # Reload the small catalog at request time so a catalog update never
        # leaves the media proxy serving a stale playback URL.
        row = read_catalog().get(paths[0])
        if row is None:
            self.send_json(404, {"error": "video_path_is_not_in_the_catalog"})
            return
        if kind == "media":
            cached_asset = cached_media_path(paths[0])
            # A cached source original must not override a catalog entry that
            # deliberately points at a browser-compatible derivative.
            cache_matches_catalog = isinstance(row.get("videoUrl"), str) and row["videoUrl"].endswith(f"/{paths[0]}")
            if cached_asset is not None and cache_matches_catalog:
                self.serve_cached_media(cached_asset, send_body=send_body)
                return
        url = row.get("videoUrl") if kind == "media" else row.get("initialFrameUrl")
        if not isinstance(url, str) or not url:
            self.send_json(404, {"error": "preview_is_not_available"})
            return
        self.proxy_asset(url, send_body=send_body)

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        path = parsed.path
        if path == "/api/health":
            self.send_json(200 if HF_TOKEN else 503, {"ok": bool(HF_TOKEN), "catalog_videos": len(CATALOG), "dataset_repo": DATASET_REPO})
            return
        if path == "/media":
            self.proxy_catalog_asset("media", parse_qs(parsed.query, keep_blank_values=True))
            return
        if path == "/preview":
            self.proxy_catalog_asset("preview", parse_qs(parsed.query, keep_blank_values=True))
            return
        if path == "/":
            self.path = "/review.html"
        super().do_GET()

    def do_HEAD(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/media":
            self.proxy_catalog_asset("media", parse_qs(parsed.query, keep_blank_values=True), send_body=False)
            return
        if parsed.path == "/preview":
            self.proxy_catalog_asset("preview", parse_qs(parsed.query, keep_blank_values=True), send_body=False)
            return
        if parsed.path == "/":
            self.path = "/review.html"
        super().do_HEAD()

    def do_POST(self) -> None:
        if urlparse(self.path).path != "/api/submissions":
            self.send_json(404, {"error": "not_found"})
            return
        if not HF_TOKEN:
            self.send_json(503, {"error": "submission_service_not_configured"})
            return
        client_ip = self.client_address[0]
        now = datetime.now(timezone.utc).timestamp()
        with rate_lock:
            previous = last_submission_by_ip.get(client_ip, 0.0)
            if now - previous < MIN_SUBMISSION_INTERVAL:
                self.send_json(429, {"error": "please_wait_before_submitting_again"})
                return
            last_submission_by_ip[client_ip] = now

        try:
            content_length = int(self.headers.get("Content-Length", "0"))
            if content_length <= 0 or content_length > MAX_BODY_BYTES:
                raise ValueError("invalid_request_size")
            payload = json.loads(self.rfile.read(content_length).decode("utf-8"))
            record = validate_submission(payload)
            submission_id = uuid.uuid4().hex
            created_at = utc_now()
            record.update({
                "schema_version": "worldmodelbench-web-review-v1",
                "submission_id": submission_id,
                "submitted_at": created_at,
                "source": "public_review_space",
            })
            path_in_repo = f"web-review/submissions/{created_at[:10]}/{submission_id}.json"
            content = json.dumps(record, ensure_ascii=False, indent=2).encode("utf-8")
            HfApi(token=HF_TOKEN).create_commit(
                repo_id=DATASET_REPO,
                repo_type="dataset",
                operations=[CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(content))],
                commit_message=f"Add web review annotation {submission_id[:12]}",
            )
            self.send_json(201, {"ok": True, "submission_id": submission_id, "submitted_at": created_at})
        except ValueError as error:
            self.send_json(400, {"error": str(error)})
        except Exception:
            self.send_json(502, {"error": "huggingface_submission_failed"})


def main() -> None:
    port = int(os.environ.get("PORT", "7860"))
    server = ThreadingHTTPServer(("0.0.0.0", port), ReviewHandler)
    print(f"WorldModelBench review service listening on {port}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()