app.py

"""Common Crawl WARC Repackaging — Hugging Face Space (Gradio).

A simple form to repackage the Common Crawl records for a set of hostnames/domains
across selected crawls into a new WARC archive in your HF bucket. Heavy lifting runs
on HF Jobs via `cdxt repackage`; this app only collects input, estimates cost, and
launches/monitors the jobs.
"""
from __future__ import annotations

import time
from typing import List, Optional

import gradio as gr

from src import buckets, config
from src.estimate import format_cost_report, format_estimate, parse_estimate
from src.job_runner import TERMINAL_STAGES, make_runner
from src.jobs import RepackageRequest, bucket_uris, bucket_web_url

RUNNER = make_runner()
IS_PODMAN = config.EXECUTOR == "podman"
POLL_SECONDS = 2 if IS_PODMAN else 5
# Safety cap on how long the UI keeps polling a job (the job's own --timeout still applies).
MAX_POLL_SECONDS = 20 * 60


def _token_str(oauth_token) -> Optional[str]:
    if IS_PODMAN:
        return None
    return getattr(oauth_token, "token", None)


def _build_request(
    domains_text: str,
    hostnames_text: str,
    languages: List[str],
    crawls: List[str],
    target_bucket: str,
    target_path: str,
    name: str,
    flavor: str,
    profile,
) -> tuple[Optional[RepackageRequest], str]:
    domains = buckets.parse_targets(domains_text)
    hostnames = buckets.parse_targets(hostnames_text)
    languages = list(languages or [])

    ok, msg = buckets.validate_targets(hostnames, domains)
    if not ok:
        return None, f"❌ {msg}"
    ok, msg = buckets.validate_languages(languages)
    if not ok:
        return None, f"❌ {msg}"
    ok, msg = buckets.validate_crawls(crawls)
    warning = msg if (ok and msg) else ""
    if not ok:
        return None, f"❌ {msg}"

    creator = getattr(profile, "username", "") or ""
    req = RepackageRequest(
        hostnames=hostnames,
        domains=domains,
        languages=languages,
        crawls=list(crawls),
        target_bucket=target_bucket.strip(),
        target_path=target_path.strip(),
        name=(name.strip() or "repackaged"),
        flavor=flavor,
        creator=creator,
    )
    return req, warning


def refresh_crawls(oauth_token: gr.OAuthToken | None = None):
    crawls = buckets.list_available_crawls(_token_str(oauth_token))
    return gr.update(choices=crawls, value=[])


def refresh_buckets(oauth_token: gr.OAuthToken | None = None):
    """Populate the target-bucket dropdown with buckets the user can write to."""
    if IS_PODMAN:
        return gr.update(choices=["local/out"], value="local/out")
    return gr.update(choices=buckets.list_writable_buckets(_token_str(oauth_token)))


def _poll(job_id: str, token: Optional[str], header: str):
    """Generator yielding (status_markdown, logs_text) until the job is terminal."""
    started = time.monotonic()
    url = RUNNER.job_url(job_id)
    link = f" — [open job]({url})" if url else ""
    while True:
        stage = RUNNER.status(job_id, token)
        logs = RUNNER.logs(job_id, token)
        yield (f"{header}\n\n**Job** `{job_id}`{link}\n\n**Status:** `{stage}`",
               "\n".join(logs[-400:]))
        if stage in TERMINAL_STAGES:
            return
        if time.monotonic() - started > MAX_POLL_SECONDS:
            yield (f"{header}\n\n**Job** `{job_id}`{link}\n\n**Status:** `{stage}` "
                   "(stopped polling; check the job page)", "\n".join(logs[-400:]))
            return
        time.sleep(POLL_SECONDS)


def estimate_handler(
    domains_text, hostnames_text, languages, crawls, target_bucket, target_path, name, flavor,
    profile: gr.OAuthProfile | None = None,
    oauth_token: gr.OAuthToken | None = None,
):
    # Outputs: (estimate_md, logs, launch_btn, est_records, active_job_id, cancel_btn)
    disabled = gr.update(interactive=False)
    enabled = gr.update(interactive=True)
    cancel_on = gr.update(interactive=True)
    cancel_off = gr.update(interactive=False)

    if not IS_PODMAN and _token_str(oauth_token) is None:
        yield "❌ Please sign in with Hugging Face first.", "", disabled, 0, "", cancel_off
        return

    req, msg = _build_request(
        domains_text, hostnames_text, languages, crawls, target_bucket, target_path,
        name, flavor, profile,
    )
    if req is None:
        yield msg, "", disabled, 0, "", cancel_off
        return

    token = _token_str(oauth_token)
    if not IS_PODMAN:
        ok, vmsg = buckets.validate_target_bucket(req.target_bucket, req.target_path, token)
        if not ok:
            yield f"❌ {vmsg}", "", disabled, 0, "", cancel_off
            return

    header = "### ⏳ Estimating…" + (f"\n\n{msg}" if msg else "")
    try:
        job_id = RUNNER.estimate(req, token)
    except Exception as e:  # noqa: BLE001
        yield f"❌ Failed to start estimate job: {type(e).__name__}: {e}", "", disabled, 0, "", cancel_off
        return

    final_logs: List[str] = []
    for status_md, logs_text in _poll(job_id, token, header):
        final_logs = logs_text.splitlines()
        # job_id active -> cancellable; enable the Cancel button while it runs.
        yield status_md, logs_text, disabled, 0, job_id, cancel_on

    est = parse_estimate(final_logs, req.flavor)
    if est is None:
        err = next((ln for ln in reversed(final_logs) if "ERROR" in ln), "")
        detail = f"\n\n`{err.strip()}`" if err else " See logs below."
        yield (f"❌ Estimate job failed — no ESTIMATE summary was produced.{detail}",
               "\n".join(final_logs[-400:]), disabled, 0, "", cancel_off)
        return

    summary = format_estimate(est)
    if msg:
        summary += f"\n\n{msg}"
    report = RUNNER.report(job_id, token)
    if report:
        summary += "\n\n" + format_cost_report(*report)
    yield (summary, "\n".join(final_logs[-400:]),
           (enabled if est.n_records > 0 else disabled), est.n_records, "", cancel_off)


def launch_handler(
    domains_text, hostnames_text, languages, crawls, target_bucket, target_path, name, flavor,
    est_records,
    profile: gr.OAuthProfile | None = None,
    oauth_token: gr.OAuthToken | None = None,
):
    # Outputs: (launch_status, logs, active_job_id, cancel_btn)
    cancel_on = gr.update(interactive=True)
    cancel_off = gr.update(interactive=False)

    if not IS_PODMAN and _token_str(oauth_token) is None:
        yield "❌ Please sign in with Hugging Face first.", "", "", cancel_off
        return

    req, msg = _build_request(
        domains_text, hostnames_text, languages, crawls, target_bucket, target_path,
        name, flavor, profile,
    )
    if req is None:
        yield msg, "", "", cancel_off
        return

    token = _token_str(oauth_token)
    header = "### 🚀 Repackaging"
    try:
        job_id = RUNNER.fetch(req, token, n_records=int(est_records or 0))
    except Exception as e:  # noqa: BLE001
        yield f"❌ Failed to start repackage job: {type(e).__name__}: {e}", "", "", cancel_off
        return

    last_status, last_logs = header, ""
    for status_md, logs_text in _poll(job_id, token, header):
        last_status, last_logs = status_md, logs_text
        # job_id active -> cancellable; enable the Cancel button while it runs.
        yield status_md, logs_text, job_id, cancel_on

    final_stage = RUNNER.status(job_id, token)
    if final_stage == "COMPLETED":
        ranges_uri, output_uri = bucket_uris(req)
        bucket_link = bucket_web_url(req.target_bucket, req.normalized_path())
        done = (
            "### ✅ Repackaging complete\n"
            f"Output written to `{output_uri}`\n\n"
            f"**[📂 Open in your bucket]({bucket_link})**\n\n"
            f"Range-jobs CSV: `{ranges_uri}`"
        )
        report = RUNNER.report(job_id, token)
        if report:
            done += "\n\n" + format_cost_report(*report)
        yield done, last_logs, "", cancel_off
    else:
        yield (f"{last_status}\n\n### ❌ Job did not complete successfully — see logs.",
               last_logs, "", cancel_off)


def cancel_handler(active_job, oauth_token: gr.OAuthToken | None = None):
    # Cancels ONLY the job this session launched (active_job holds that specific id),
    # never any other job the user may be running.
    off = gr.update(interactive=False)
    if not active_job:
        return "No running job to cancel.", off
    token = _token_str(oauth_token)
    try:
        RUNNER.cancel(active_job, token)
        return f"🛑 Cancel requested for job `{active_job}`.", off
    except Exception as e:  # noqa: BLE001
        return f"❌ Failed to cancel job `{active_job}`: {type(e).__name__}: {e}", off


def build_demo() -> gr.Blocks:
    # Use the static fallback at startup (no network); the Refresh button does a
    # live, authenticated listing.
    initial_crawls = list(config.FALLBACK_CRAWLS)
    with gr.Blocks(title="CC WARC Repackager") as demo:
        gr.Markdown(
            "# Common Crawl WARC Repackager\n"
            "Repackage the Web page records for your hostnames/domains across selected "
            "Common Crawl archives into a new WARC archive in your Hugging Face bucket. Compute runs on "
            "[HF Jobs](https://huggingface.co/docs/hub/jobs) data is read "
            "from the [Common Crawl HF bucket](https://huggingface.co/buckets/commoncrawl/commoncrawl)."
        )

        if IS_PODMAN:
            gr.Markdown("> 🧪 **Local podman mode** — jobs run in a local container against fixtures.")
        else:
            gr.Markdown(
                "Jobs are available to any user or organization with a [positive credit balance](https://huggingface.co/settings/billing)."
            )
            gr.LoginButton()

        gr.Markdown("## 1. Filter target")
        with gr.Row():
            domains_text = gr.Textbox(
                label="Registered domains",
                placeholder="example.com, example.org",
                info="Matches the domain and all its subdomains (url_host_registered_domain).",
            )
            hostnames_text = gr.Textbox(
                label="Exact hostnames",
                placeholder="www.example.com, www.example.org",
                info="Matches exact hosts only (url_host_name). Combine with domains.",
            )
        with gr.Row():
            languages = gr.Dropdown(
                choices=config.language_choices(),
                value=[],
                multiselect=True,
                label="Content languages (optional)",
                info="ISO-639-3 codes (content_languages). Keeps records whose primary "
                     "detected language is one of these. Leave empty for all languages.",
            )

        gr.Markdown("## 2. Target crawls")
        crawls = gr.CheckboxGroup(
            choices=initial_crawls, value=initial_crawls[:1], label="Crawls to scan",
            info=f"Selecting > {config.COST_GUARD_MAX_CRAWLS} crawls makes the index scan large.",
        )
        refresh_btn = gr.Button("🔄 Refresh crawl list", size="sm")

        gr.Markdown("## 3. Output target")
        gr.Markdown(
            "Pick one of your **existing** buckets (the dropdown lists buckets you can write "
            "to). Create a new one at https://huggingface.co/new-bucket, then click Refresh. "
            "Dataset repos are not supported as a target."
        )
        with gr.Row(equal_height=True):
            target_bucket = gr.Dropdown(
                choices=(["local/out"] if IS_PODMAN else []),
                value=("local/out" if IS_PODMAN else None),
                label="Target bucket",
                allow_custom_value=True,
                info="Buckets you can write to (yours + write-access orgs).",
            )
            target_path = gr.Textbox(
                label="Path within bucket", placeholder="exports/run1", value="",
                info="Sub-folder for the output; created if missing. Leave blank for the bucket root.",
            )
        refresh_buckets_btn = gr.Button("🔄 Refresh buckets", size="sm")
        with gr.Row():
            name = gr.Textbox(label="Output WARC name", value="repackaged")
            flavor = gr.Dropdown(
                choices=list(config.FLAVOR_HOURLY_USD.keys()), value=config.DEFAULT_FLAVOR,
                label="Hardware flavor",
            )

        gr.Markdown("## 4. Estimate & launch")
        with gr.Row():
            estimate_btn = gr.Button("📊 Estimate cost", variant="secondary")
            launch_btn = gr.Button("🚀 Launch repackage", variant="primary", interactive=False)
            cancel_btn = gr.Button("🛑 Cancel job", variant="stop", interactive=False)

        estimate_md = gr.Markdown("")
        launch_status = gr.Markdown("")
        logs_box = gr.Textbox(label="Job logs", lines=16, max_lines=16, autoscroll=True)
        est_records = gr.State(0)   # n_records from the estimate, sizes fetch --processes
        active_job = gr.State("")   # id of the currently running job, for cancellation

        form_inputs = [
            domains_text, hostnames_text, languages, crawls, target_bucket, target_path,
            name, flavor,
        ]

        refresh_btn.click(refresh_crawls, outputs=[crawls], api_name="refresh_crawls")
        refresh_buckets_btn.click(refresh_buckets, outputs=[target_bucket], api_name="refresh_buckets")
        # Auto-populate the bucket dropdown on load (after OAuth the token is available).
        demo.load(refresh_buckets, outputs=[target_bucket])
        estimate_event = estimate_btn.click(
            estimate_handler, inputs=form_inputs,
            outputs=[estimate_md, logs_box, launch_btn, est_records, active_job, cancel_btn],
            api_name="estimate",
        )
        launch_event = launch_btn.click(
            launch_handler, inputs=form_inputs + [est_records],
            outputs=[launch_status, logs_box, active_job, cancel_btn], api_name="launch",
        )
        # Cancel stops the streaming estimate/launch generators AND cancels that job
        # (only the one this session launched — active_job holds its id).
        cancel_btn.click(
            cancel_handler, inputs=[active_job], outputs=[launch_status, cancel_btn],
            cancels=[estimate_event, launch_event], api_name="cancel",
        )

    return demo


if __name__ == "__main__":
    build_demo().queue().launch()