main.py

import os
import logging
from typing import Any

from fastapi import FastAPI, Header, HTTPException, Query, Request
from huggingface_hub import HfApi


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

HF_TOKEN = os.environ["HF_TOKEN"]
WEBHOOK_SECRET = os.environ["WEBHOOK_SECRET"]

api = HfApi(token=HF_TOKEN)


def _normalize_prefix(value: str) -> str:
    value = value.strip().strip("/")
    return f"{value}/" if value else ""


def _load_allowed_prefixes() -> list[str]:
    raw = os.getenv("ALLOWED_PREFIXES", "").strip()
    if raw:
        prefixes = sorted(
            {
                _normalize_prefix(item)
                for item in raw.split(",")
                if _normalize_prefix(item)
            }
        )
        logger.info("Using ALLOWED_PREFIXES override: %s", prefixes)
        return prefixes

    me = api.whoami()
    prefixes = {_normalize_prefix(me["name"])}

    for org in me.get("orgs", []):
        if isinstance(org, dict):
            name = org.get("name", "")
        else:
            name = str(org)
        normalized = _normalize_prefix(name)
        if normalized:
            prefixes.add(normalized)

    prefixes = sorted(prefixes)
    logger.info("Derived allowed prefixes from whoami(): %s", prefixes)
    return prefixes


ALLOWED_PREFIXES = _load_allowed_prefixes()


def _namespace_of_repo_id(repo_id: str) -> str:
    parts = repo_id.strip().split("/", 1)
    if len(parts) != 2 or not parts[0] or not parts[1]:
        raise ValueError("repo_id must be in 'namespace/name' format")
    return parts[0]


def _is_allowed_space(repo_id: str) -> bool:
    try:
        namespace = _namespace_of_repo_id(repo_id)
    except ValueError:
        return False
    return _normalize_prefix(namespace) in ALLOWED_PREFIXES


def _get_webhook_secret(
    x_webhook_secret: str | None,
    secret_query: str | None,
) -> str | None:
    return x_webhook_secret or secret_query


@app.get("/healthz")
async def healthz() -> dict[str, Any]:
    return {
        "ok": True,
        "allowed_prefixes": ALLOWED_PREFIXES,
    }


@app.post("/webhook")
async def handle_webhook(
    request: Request,
    space: str = Query(..., description="Target Space repo_id, e.g. user/my-space"),
    expected_repo: str | None = Query(
        default=None,
        description="Optional source repo_id to require for this webhook",
    ),
    secret: str | None = Query(
        default=None,
        description="Optional webhook secret in query params if header is unavailable",
    ),
    x_webhook_secret: str | None = Header(default=None),
) -> dict[str, Any]:
    supplied_secret = _get_webhook_secret(x_webhook_secret, secret)
    if supplied_secret != WEBHOOK_SECRET:
        raise HTTPException(status_code=403, detail="Bad secret")

    if not _is_allowed_space(space):
        raise HTTPException(
            status_code=403,
            detail=f"Target Space namespace is not allowed: {space}",
        )

    try:
        payload = await request.json()
    except Exception as exc:
        raise HTTPException(status_code=400, detail="Invalid JSON payload") from exc

    repo = payload.get("repo", {})
    event = payload.get("event", {})

    repo_name = repo.get("name")
    repo_type = repo.get("type")
    event_action = event.get("action")
    event_scope = event.get("scope")

    if repo_type != "model":
        return {
            "status": "ignored",
            "reason": "repo is not a model",
            "source_repo": repo_name,
            "space": space,
        }

    if event_action != "update":
        return {
            "status": "ignored",
            "reason": "event action is not update",
            "source_repo": repo_name,
            "space": space,
        }

    if event_scope not in (None, "repo.content"):
        return {
            "status": "ignored",
            "reason": "event scope is not repo.content",
            "source_repo": repo_name,
            "space": space,
        }

    if expected_repo and repo_name != expected_repo:
        return {
            "status": "ignored",
            "reason": "source repo does not match expected_repo",
            "source_repo": repo_name,
            "expected_repo": expected_repo,
            "space": space,
        }

    logger.info(
        "Restarting Space '%s' due to update on source repo '%s'",
        space,
        repo_name,
    )
    api.restart_space(repo_id=space)

    return {
        "status": "restarted",
        "space": space,
        "source_repo": repo_name,
        "allowed_prefixes": ALLOWED_PREFIXES,
    }