Blame for server/main.py

"""FastAPI backend for the Common Crawl WARC Repackager wizard UI.

This is a thin HTTP + SSE shell over the existing, executor-agnostic business logic in
``src/`` — it does NOT reimplement any job logic:

  * ``src.job_runner.make_runner`` picks the real HF-Jobs runner or the local podman
    runner from ``CC_REPACKAGE_EXECUTOR`` (same as the Gradio app).
  * ``src.jobs`` builds the job specs / bucket URIs, ``src.buckets`` validates input and
    lists crawls/buckets, ``src.estimate`` parses the ``ESTIMATE`` line and does the cost
    math, ``src.examples`` provides the precomputed demo presets + simulated log lines.

Endpoints: static wizard at ``/``; JSON at ``/api/*``; per-job log streaming over SSE at
``/api/jobs/{id}/events`` and ``/api/examples/{key}/events``. OAuth (for the deployed
Space) is wired only when the Space injects ``OAUTH_CLIENT_ID`` — locally the podman
executor runs as a guest, exactly like the Gradio app's ``IS_PODMAN`` path.
"""
from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Dict, Optional

from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from src import buckets, config, examples
from src.estimate import Estimate, actual_cost_usd, parse_estimate
from src.job_runner import TERMINAL_STAGES, make_runner
from src.jobs import RepackageRequest, bucket_file_url, bucket_uris, bucket_web_url

STATIC_DIR = Path(__file__).parent / "static"
IS_PODMAN = config.EXECUTOR == "podman"
POLL_SECONDS = 2 if IS_PODMAN else 5
EXAMPLE_LINE_DELAY = 0.35

RUNNER = make_runner()
# job_id -> {kind, req, token}; lets the SSE stream build the right terminal payload.
JOBS: Dict[str, dict] = {}

app = FastAPI(title="CC WARC Repackager")


# --- auth (guest locally; real OAuth on the Space) ---------------------------
def _session(request: Request):
    # request.session only exists when SessionMiddleware is installed (OAuth wired).
    # `hasattr` is unsafe here — the property raises AssertionError, not AttributeError.
    return request.scope.get("session")


def token_for(request: Request) -> Optional[str]:
    if IS_PODMAN:
        return None
    sess = _session(request)
    return sess.get("token") if sess else None


def user_for(request: Request) -> Optional[str]:
    if IS_PODMAN:
        return None
    sess = _session(request)
    return sess.get("user") if sess else None


# --- serialization helpers ---------------------------------------------------
def _estimate_json(est: Estimate) -> dict:
    return {
        "n_records": est.n_records,
        "total_bytes": est.total_bytes,
        "total_gib": round(est.total_gib, 4),
        "n_output_files": est.n_output_files,
        "rec_per_s": est.rec_per_s,
        "hourly_usd": est.hourly_usd,
        "duration_seconds": round(est.duration_seconds, 1),
        "cost_usd": round(est.cost_usd, 4),
        "flavor": est.flavor,
    }


def _example_json(ex: examples.Example) -> dict:
    return {
        "key": ex.key,
        "title": ex.title,
        "description": ex.description,
        "kind": "sql" if ex.sql_where else "form",
        "n_records": ex.n_records,
        "total_bytes": ex.total_bytes,
        "download_url": bucket_file_url(ex.bucket, ex.path, ex.files[0]) if ex.files else "",
        "folder_url": bucket_web_url(ex.bucket, ex.path),
        "file": ex.files[0] if ex.files else "",
        "bucket": ex.bucket,
        "path": ex.path,
        "form": {
            "domains": ", ".join(ex.domains),
            "hostnames": ", ".join(ex.hostnames),
            "languages": list(ex.languages),
            "sql_where": ex.sql_where,
            "max_records": ex.max_records,
            "crawls": list(ex.crawls),
            "name": ex.name,
        },
    }


def _sse(payload: dict) -> str:
    return f"data: {json.dumps(payload)}\n\n"


def clean_log(text: str) -> str:
    """Trim the local podman runner's echoed job script from the log for readability.

    The local runner writes the full `podman run ... <heredoc>` command (which inlines
    index_query.py) ahead of the container's own stdout. Real HF Jobs logs contain only
    the container stdout, so this is a no-op there (no marker / no `$ podman` line)."""
    marker = "CC_PY_EOF"
    if marker in text:                       # estimate job: keep output after the heredoc
        text = text.rsplit(marker, 1)[-1]
    lines = [ln for ln in text.splitlines() if not ln.startswith("$ podman run")]
    return "\n".join(lines).strip()


# --- request building (mirrors the Gradio app's _build_request) --------------
def build_request(p: dict) -> tuple[Optional[RepackageRequest], Optional[str], str]:
    """Return (request, error, warning). Validates via the shared src.buckets helpers."""
    domains = buckets.parse_targets(p.get("domains", ""))
    hostnames = buckets.parse_targets(p.get("hostnames", ""))
    languages = list(p.get("languages") or [])
    sql_where = (p.get("sql_where") or "").strip()
    warnings = []

    ok, msg, max_records = buckets.validate_max_records(p.get("max_records"))
    if not ok:
        return None, msg, ""

    if sql_where:
        ok, msg = buckets.validate_sql_where(sql_where)
        if not ok:
            return None, msg, ""
        if hostnames or domains or languages:
            warnings.append("Using the advanced SQL filter — the domain / hostname / language fields are ignored.")
        hostnames, domains, languages = [], [], []
    else:
        ok, msg = buckets.validate_languages(languages)
        if not ok:
            return None, msg, ""
        ok, msg = buckets.validate_targets(hostnames, domains, languages)
        if not ok:
            return None, msg, ""

    crawls = list(p.get("crawls") or [])
    ok, msg = buckets.validate_crawls(crawls)
    if not ok:
        return None, msg, ""
    if msg:
        warnings.append(msg)

    req = RepackageRequest(
        hostnames=hostnames, domains=domains, languages=languages, sql_where=sql_where,
        max_records=max_records, crawls=crawls,
        target_bucket=(p.get("target_bucket") or "").strip(),
        target_path=(p.get("target_path") or "").strip(),
        name=(p.get("name") or "").strip() or "repackaged",
        flavor=p.get("flavor") or config.DEFAULT_FLAVOR,
        creator=p.get("creator") or "",
    )
    return req, None, "\n".join(warnings)


# --- API ---------------------------------------------------------------------
@app.get("/api/config")
def api_config(request: Request):
    return {
        "is_podman": IS_PODMAN,
        "signed_in": bool(token_for(request)) or IS_PODMAN,
        "user": user_for(request),
        "default_flavor": config.DEFAULT_FLAVOR,
        "flavors": {
            f: {"usd": config.FLAVOR_HOURLY_USD[f], "rps": config.FLAVOR_REC_PER_S.get(f, 0),
                "vcpus": config.FLAVOR_VCPUS.get(f, 0)}
            for f in config.FLAVOR_HOURLY_USD
        },
        "crawls": list(config.FALLBACK_CRAWLS),
        "languages": [[n, c] for n, c in config.CONTENT_LANGUAGES],
        "cost_guard_max_crawls": config.COST_GUARD_MAX_CRAWLS,
        "examples": [_example_json(ex) for ex in examples.EXAMPLES],
    }


@app.get("/api/crawls")
def api_crawls(request: Request):
    if IS_PODMAN:
        return {"crawls": list(config.FALLBACK_CRAWLS)}
    return {"crawls": buckets.list_available_crawls(token_for(request))}


@app.get("/api/buckets")
def api_buckets(request: Request):
    if IS_PODMAN:
        return {"buckets": ["local/out"]}
    return {"buckets": buckets.list_writable_buckets(token_for(request))}


@app.post("/api/validate")
async def api_validate(request: Request):
    req, err, warning = build_request(await request.json())
    if err:
        return {"ok": False, "error": err}
    return {"ok": True, "warning": warning}


@app.post("/api/estimate")
async def api_estimate(request: Request):
    p = await request.json()
    token = token_for(request)
    if not IS_PODMAN and token is None:
        return JSONResponse({"error": "Please sign in with Hugging Face first."}, status_code=401)
    req, err, warning = build_request(p)
    if err:
        return JSONResponse({"error": err}, status_code=400)
    if not IS_PODMAN:
        ok, vmsg = buckets.validate_target_bucket(req.target_bucket, req.target_path, token)
        if not ok:
            return JSONResponse({"error": vmsg}, status_code=400)
    try:
        job_id = RUNNER.estimate(req, token)
    except Exception as e:  # noqa: BLE001
        return JSONResponse({"error": f"{type(e).__name__}: {e}"}, status_code=500)
    JOBS[job_id] = {"kind": "estimate", "req": req, "token": token}
    return {"job_id": job_id, "warning": warning}


@app.post("/api/launch")
async def api_launch(request: Request):
    p = await request.json()
    token = token_for(request)
    if not IS_PODMAN and token is None:
        return JSONResponse({"error": "Please sign in with Hugging Face first."}, status_code=401)
    req, err, warning = build_request(p)
    if err:
        return JSONResponse({"error": err}, status_code=400)
    try:
        job_id = RUNNER.fetch(req, token, n_records=int(p.get("n_records") or 0))
    except Exception as e:  # noqa: BLE001
        return JSONResponse({"error": f"{type(e).__name__}: {e}"}, status_code=500)
    JOBS[job_id] = {"kind": "fetch", "req": req, "token": token}
    return {"job_id": job_id, "warning": warning}


@app.post("/api/jobs/{job_id}/cancel")
def api_cancel(job_id: str, request: Request):
    info = JOBS.get(job_id)
    token = info["token"] if info else token_for(request)
    try:
        RUNNER.cancel(job_id, token)
        return {"ok": True}
    except Exception as e:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": f"{type(e).__name__}: {e}"}, status_code=500)


@app.get("/api/jobs/{job_id}/events")
def api_job_events(job_id: str):
    info = JOBS.get(job_id)
    if not info:
        return JSONResponse({"error": "unknown job"}, status_code=404)
    req, token, kind = info["req"], info["token"], info["kind"]

    def gen():
        url = RUNNER.job_url(job_id)
        if url:
            yield _sse({"type": "url", "url": url})
        started = time.monotonic()
        final_logs = []
        while True:
            stage = RUNNER.status(job_id, token)
            logs = RUNNER.logs(job_id, token)
            final_logs = logs
            yield _sse({"type": "tick", "stage": stage, "log": clean_log("\n".join(logs[-600:]))})
            if stage in TERMINAL_STAGES or time.monotonic() - started > 20 * 60:
                break
            time.sleep(POLL_SECONDS)

        report = RUNNER.report(job_id, token)
        done = {"type": "done", "stage": stage}
        if report:
            secs, flavor = report
            done["report"] = {"seconds": secs, "flavor": flavor,
                              "cost_usd": round(actual_cost_usd(secs, flavor), 4)}
        if kind == "estimate":
            est = parse_estimate(final_logs, req.flavor)
            if est is not None:
                done["estimate"] = _estimate_json(est)
            else:
                cleaned = clean_log("\n".join(final_logs)).splitlines()
                err = next((ln for ln in reversed(cleaned) if "ERROR" in ln or "Error" in ln), "")
                done["error"] = err.strip() or "No ESTIMATE summary was produced."
        else:
            ranges_uri, output_uri = bucket_uris(req)
            done["result"] = {
                "ranges_uri": ranges_uri, "output_uri": output_uri,
                "bucket_url": bucket_web_url(req.target_bucket, req.normalized_path()),
                "completed": stage == "COMPLETED",
            }
        yield _sse(done)

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.get("/api/examples/{key}/events")
def api_example_events(key: str, phase: Optional[str] = None):
    """Stream a precomputed example's simulated logs. `phase` selects a single step of
    the tutorial: "estimate" (index logs + cost) or "fetch" (repackage logs + download).
    Omitting it streams both back to back (legacy)."""
    ex = examples.EXAMPLES_BY_KEY.get(key)
    if not ex:
        return JSONResponse({"error": "unknown example"}, status_code=404)

    def gen():
        if phase in (None, "estimate"):
            for line in examples.estimate_log_lines(ex):
                time.sleep(EXAMPLE_LINE_DELAY)
                yield _sse({"type": "log", "line": line})
            yield _sse({"type": "estimate", "estimate": _estimate_json(
                Estimate(ex.n_records, ex.total_bytes, config.DEFAULT_FLAVOR))})
            if phase == "estimate":
                yield _sse({"type": "done"})
                return
            time.sleep(0.3)
        if phase in (None, "fetch"):
            for line in examples.fetch_log_lines(ex):
                time.sleep(EXAMPLE_LINE_DELAY)
                yield _sse({"type": "log", "line": line})
            yield _sse({"type": "done", "result": {
                "demo": True, "download_url": bucket_file_url(ex.bucket, ex.path, ex.files[0]),
                "folder_url": bucket_web_url(ex.bucket, ex.path), "file": ex.files[0],
                "n_records": ex.n_records, "total_bytes": ex.total_bytes}})

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


# --- OAuth (only active on the Space, when HF injects the client id) ----------
def _redirect_uri(request: Request) -> str:
    """The OAuth callback URL. On a Space, HF registers https://<SPACE_HOST>/auth/callback,
    and requests arrive behind a proxy, so prefer SPACE_HOST over the request's own host."""
    host = os.environ.get("SPACE_HOST")
    if host:
        return f"https://{host}/auth/callback"
    return str(request.url_for("auth"))


def _wire_oauth():
    client_id = os.environ.get("OAUTH_CLIENT_ID")
    if not client_id:
        return  # local / podman: guest mode, no OAuth
    try:
        from authlib.integrations.starlette_client import OAuth
        from starlette.middleware.sessions import SessionMiddleware
        from starlette.responses import RedirectResponse
    except Exception:  # noqa: BLE001
        return
    app.add_middleware(
        SessionMiddleware,
        secret_key=os.environ.get("SESSION_SECRET") or os.environ.get("OAUTH_CLIENT_SECRET") or os.urandom(24).hex(),
        same_site="lax", https_only=True,
    )
    # HF login needs openid+profile for userinfo, plus the app scopes (write-repos, jobs).
    scopes = os.environ.get("OAUTH_SCOPES", "write-repos jobs")
    for s in ("profile", "openid"):
        if s not in scopes.split():
            scopes = s + " " + scopes
    oauth = OAuth()
    oauth.register(
        name="hf",
        client_id=client_id,
        client_secret=os.environ.get("OAUTH_CLIENT_SECRET"),
        server_metadata_url=os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co") + "/.well-known/openid-configuration",
        client_kwargs={"scope": scopes},
    )

    @app.get("/login")
    async def login(request: Request):
        return await oauth.hf.authorize_redirect(request, _redirect_uri(request))

    @app.get("/auth/callback", name="auth")
    async def auth(request: Request):
        try:
            tok = await oauth.hf.authorize_access_token(request)
        except Exception as e:  # noqa: BLE001
            return JSONResponse({"error": f"OAuth failed: {type(e).__name__}: {e}"}, status_code=400)
        request.session["token"] = tok.get("access_token")
        ui = tok.get("userinfo") or {}
        request.session["user"] = ui.get("preferred_username") or ui.get("name")
        return RedirectResponse(url="/")

    @app.get("/logout")
    async def logout(request: Request):
        request.session.clear()
        return RedirectResponse(url="/")


_wire_oauth()


# --- static (mounted at /static; "/" serves the app shell) -------------------
@app.get("/")
def index():
    return FileResponse(STATIC_DIR / "index.html")


app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")