backend/main.py

from fastapi import FastAPI
from fastapi import HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import HfApi
from huggingface_hub import hf_hub_download
from pydantic import BaseModel
from typing import Optional
import csv
from datetime import datetime
from datetime import timezone
import io
import json
import os
from pathlib import Path
import re
from uuid import uuid4

app = FastAPI()

frontend_origins = os.getenv("FRONTEND_ORIGINS", "*")
allow_origins = (
    ["*"]
    if frontend_origins == "*"
    else [origin.strip() for origin in frontend_origins.split(",") if origin.strip()]
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=allow_origins,
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

HF_DATASET_REPO_ID = os.getenv(
    "HF_ANNOTATIONS_DATASET",
    "adinayak/t2av_eval_annotations",
)
HF_TOKEN = os.getenv("HF_TOKEN")
HF_DATASET_PRIVATE = os.getenv("HF_ANNOTATIONS_DATASET_PRIVATE", "true").lower() in {
    "1",
    "true",
    "yes",
}
ANNOTATIONS_DIR = "annotations"
api = HfApi()


class AnnotationPayload(BaseModel):
    video_id: str
    annotations: dict
    user: Optional[str] = "anonymous"


@app.on_event("startup")
def startup():
    ensure_dataset_configured()


@app.get("/health")
def root():
    return {
        "status": "running",
        "annotation_store": "huggingface_dataset",
        "dataset_repo": HF_DATASET_REPO_ID,
    }


def ensure_dataset_configured():
    if not HF_DATASET_REPO_ID:
        raise HTTPException(
            status_code=500,
            detail="HF_ANNOTATIONS_DATASET is not configured.",
        )


def ensure_write_token():
    if not HF_TOKEN:
        raise HTTPException(
            status_code=500,
            detail="HF_TOKEN is not configured. Add it as a Hugging Face Space secret.",
        )


def ensure_dataset_exists():
    ensure_dataset_configured()
    ensure_write_token()
    api.create_repo(
        repo_id=HF_DATASET_REPO_ID,
        repo_type="dataset",
        private=HF_DATASET_PRIVATE,
        token=HF_TOKEN,
        exist_ok=True,
    )


def slugify(value):
    slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-")
    return slug or "unknown"


def annotation_path(annotation):
    created_at = datetime.fromisoformat(annotation["created_at"].replace("Z", "+00:00"))
    timestamp = created_at.strftime("%Y%m%dT%H%M%S%fZ")
    video_id = slugify(annotation["video_id"])
    user = slugify(annotation["user"])
    return f"{ANNOTATIONS_DIR}/{timestamp}_{video_id}_{user}_{annotation['id']}.json"


def read_annotation_file(path):
    try:
        local_path = hf_hub_download(
            repo_id=HF_DATASET_REPO_ID,
            filename=path,
            repo_type="dataset",
            token=HF_TOKEN,
        )
    except Exception as exc:
        raise HTTPException(status_code=502, detail="Failed to read annotation file") from exc

    with open(local_path, encoding="utf-8") as annotation_file:
        return json.load(annotation_file)


def list_annotation_paths():
    ensure_dataset_configured()
    try:
        files = api.list_repo_files(
            repo_id=HF_DATASET_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN,
        )
    except Exception as exc:
        raise HTTPException(status_code=502, detail="Failed to list annotation dataset") from exc

    return sorted(
        [
            path
            for path in files
            if path.startswith(f"{ANNOTATIONS_DIR}/") and path.endswith(".json")
        ],
        reverse=True,
    )


def list_annotation_records():
    annotations = [read_annotation_file(path) for path in list_annotation_paths()]
    return sorted(
        annotations,
        key=lambda item: item.get("created_at", ""),
        reverse=True,
    )


@app.post("/save_annotation")
def save_annotation(payload: AnnotationPayload):
    ensure_dataset_exists()
    created_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
    annotation = {
        "id": uuid4().hex,
        "video_id": payload.video_id,
        "user": payload.user or "anonymous",
        "annotations": payload.annotations,
        "created_at": created_at,
    }
    path_in_repo = annotation_path(annotation)
    annotation_bytes = json.dumps(annotation, indent=2).encode("utf-8")

    try:
        api.upload_file(
            path_or_fileobj=io.BytesIO(annotation_bytes),
            path_in_repo=path_in_repo,
            repo_id=HF_DATASET_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN,
            commit_message=f"Add annotation {annotation['id']}",
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail="Failed to save annotation") from exc

    return {"message": "saved", "annotation": annotation, "path": path_in_repo}


@app.get("/annotations")
def list_annotations():
    return list_annotation_records()


@app.get("/annotations.csv")
def download_annotations_csv():
    annotations = list_annotation_records()
    output = io.StringIO()
    writer = csv.DictWriter(
        output,
        fieldnames=["id", "video_id", "user", "created_at", "annotations"],
    )
    writer.writeheader()
    for annotation in annotations:
        writer.writerow(
            {
                "id": annotation["id"],
                "video_id": annotation["video_id"],
                "user": annotation["user"],
                "created_at": annotation["created_at"],
                "annotations": json.dumps(annotation["annotations"]),
            }
        )
    output.seek(0)
    return StreamingResponse(
        iter([output.getvalue()]),
        media_type="text/csv",
        headers={"Content-Disposition": "attachment; filename=annotations.csv"},
    )


@app.get("/annotations/{annotation_id}")
def get_annotation(annotation_id: str):
    for annotation in list_annotation_records():
        if annotation["id"] == annotation_id:
            return annotation
    raise HTTPException(status_code=404, detail="Annotation not found")


static_dir = Path(__file__).resolve().parent.parent / "frontend" / "dist"

if static_dir.exists():
    app.mount("/", StaticFiles(directory=static_dir, html=True), name="frontend")