#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///
"""
Search Hugging Face benchmark datasets and fetch leaderboard results.

This script is designed to be pipeline-friendly:
  - search benchmark datasets by free text, alias, task, and modality
  - fetch dataset leaderboards in normalized JSON / NDJSON / table form
  - optionally read dataset ids from stdin for chaining

It uses HF_TOKEN automatically when present.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import textwrap
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Iterable


BASE_URL = "https://huggingface.co"
DEFAULT_TIMEOUT = 30

ALIASES: dict[str, list[str]] = {
    "ocr": [
        "ocr",
        "olmocr",
        "pdf",
        "image-to-text",
        "screen",
        "screenspot",
        "markdown",
        "text recognition",
    ],
    "coding": [
        "code",
        "coding",
        "software engineering",
        "programming",
        "swe",
        "terminal",
        "patch",
        "bug",
        "cuda",
    ],
    "math": [
        "math",
        "reasoning",
        "gsm8k",
        "mmlu",
        "gpqa",
        "aime",
        "hmmt",
    ],
    "retrieval": [
        "retrieval",
        "search",
        "mteb",
        "arguana",
        "bright",
    ],
    "agents": [
        "agent",
        "agents",
        "terminal",
        "screen",
        "computer use",
        "tool use",
    ],
    "asr": [
        "asr",
        "speech",
        "audio",
        "transcribe",
        "transcription",
    ],
}


class HfApiError(RuntimeError):
    pass


class FullHelpArgumentParser(argparse.ArgumentParser):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._search_parser: argparse.ArgumentParser | None = None
        self._leaderboard_parser: argparse.ArgumentParser | None = None

    def format_help(self) -> str:
        text = super().format_help()
        extra_sections: list[str] = []

        if self._search_parser is not None:
            extra_sections.append(
                "\nsearch command options:\n"
                + textwrap.indent(self._search_parser.format_help().strip(), "  ")
            )

        if self._leaderboard_parser is not None:
            extra_sections.append(
                "\nleaderboard command options:\n"
                + textwrap.indent(self._leaderboard_parser.format_help().strip(), "  ")
            )

        if extra_sections:
            text += "\n" + "\n".join(extra_sections) + "\n"
        return text


def auth_headers() -> dict[str, str]:
    token = os.getenv("HF_TOKEN")
    return {"Authorization": f"Bearer {token}"} if token else {}


def http_get_json(path: str, params: dict[str, Any] | None = None) -> Any:
    url = f"{BASE_URL}{path}"
    if params:
        pairs: list[tuple[str, str]] = []
        for key, value in params.items():
            if value is None:
                continue
            if isinstance(value, (list, tuple)):
                for item in value:
                    pairs.append((key, str(item)))
            else:
                pairs.append((key, str(value)))
        if pairs:
            url = f"{url}?{urllib.parse.urlencode(pairs)}"

    req = urllib.request.Request(url, headers=auth_headers())
    try:
        with urllib.request.urlopen(req, timeout=DEFAULT_TIMEOUT) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise HfApiError(f"{exc.code} {exc.reason} for {url}: {body[:500]}") from exc
    except urllib.error.URLError as exc:
        raise HfApiError(f"Request failed for {url}: {exc}") from exc


def shorten(text: str, width: int) -> str:
    text = " ".join((text or "").split())
    if len(text) <= width:
        return text
    return text[: max(0, width - 1)] + "…"


def first_text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        return " ".join(first_text(v) for v in value)
    if isinstance(value, dict):
        return " ".join(first_text(v) for v in value.values())
    return str(value)


def benchmark_catalog(limit: int = 500) -> list[dict[str, Any]]:
    data = http_get_json(
        "/api/datasets",
        params={"filter": "benchmark:official", "limit": limit, "full": "true"},
    )
    if not isinstance(data, list):
        raise HfApiError("Unexpected response while listing benchmark datasets")
    return data


def dataset_search_blob(dataset: dict[str, Any]) -> str:
    card = dataset.get("cardData") or {}
    parts = [
        dataset.get("id", ""),
        dataset.get("description", ""),
        first_text(dataset.get("tags")),
        first_text(card.get("pretty_name")),
        first_text(card.get("tags")),
        first_text(card.get("task_categories")),
        first_text(card.get("task_ids")),
    ]
    return " ".join(parts).lower()


def dataset_search_fields(dataset: dict[str, Any]) -> dict[str, str]:
    card = dataset.get("cardData") or {}
    return {
        "id": first_text(dataset.get("id")).lower(),
        "pretty_name": first_text(card.get("pretty_name")).lower(),
        "tags": " ".join(
            [
                first_text(dataset.get("tags")),
                first_text(card.get("tags")),
                first_text(card.get("task_categories")),
                first_text(card.get("task_ids")),
                first_text(card.get("modality")),
            ]
        ).lower(),
        "description": first_text(dataset.get("description")).lower(),
    }


def collect_prefixed_tags(dataset: dict[str, Any], prefixes: Iterable[str]) -> list[str]:
    prefixes = tuple(prefixes)
    tags = dataset.get("tags") or []
    card = dataset.get("cardData") or {}

    out: list[str] = []
    for tag in tags:
        if isinstance(tag, str) and tag.startswith(prefixes):
            out.append(tag)

    for key, prefix in (
        ("task_categories", "task_categories:"),
        ("task_ids", "task_ids:"),
        ("modality", "modality:"),
    ):
        values = card.get(key)
        if isinstance(values, list):
            for value in values:
                full_tag = f"{prefix}{value}"
                if full_tag.startswith(prefixes):
                    out.append(full_tag)

    deduped: list[str] = []
    seen: set[str] = set()
    for tag in out:
        if tag not in seen:
            deduped.append(tag)
            seen.add(tag)
    return deduped


def expand_aliases(aliases: list[str]) -> dict[str, list[str]]:
    expanded: dict[str, list[str]] = {}
    for alias in aliases:
        terms = ALIASES.get(alias.lower(), [alias])
        expanded[alias] = terms
    return expanded


def matches_term(blob: str, term: str) -> bool:
    candidate = term.lower().strip()
    if not candidate:
        return False
    if re.fullmatch(r"[a-z0-9_]+", candidate):
        return re.search(rf"(?<![a-z0-9_]){re.escape(candidate)}(?![a-z0-9_])", blob) is not None
    return candidate in blob


def score_dataset(
    dataset: dict[str, Any],
    queries: list[str],
    aliases: dict[str, list[str]],
    tasks: list[str],
    modalities: list[str],
) -> dict[str, Any]:
    blob = dataset_search_blob(dataset)
    fields = dataset_search_fields(dataset)
    task_tags = collect_prefixed_tags(dataset, ["task_categories:", "task_ids:"])
    modality_tags = collect_prefixed_tags(dataset, ["modality:"])

    score = 0
    reasons: list[str] = []

    for query in queries:
        q = query.lower().strip()
        if q and any(matches_term(value, q) for value in fields.values()):
            score += 3
            reasons.append(f"query:{query}")

    for alias_name, terms in aliases.items():
        matched_terms: list[str] = []
        alias_score = 0
        for term in terms:
            strong_match = any(
                matches_term(fields[field_name], term)
                for field_name in ("id", "pretty_name", "tags")
            )
            desc_match = matches_term(fields["description"], term)
            if strong_match:
                alias_score += 2
                matched_terms.append(term)
            elif desc_match:
                alias_score += 1
                matched_terms.append(term)
        if matched_terms:
            score += alias_score
            reasons.append(f"alias:{alias_name}=" + ",".join(matched_terms[:5]))

    lower_task_tags = [t.lower() for t in task_tags]
    for task in tasks:
        task = task.lower().strip()
        if not task:
            continue
        exact = [
            tag
            for tag in lower_task_tags
            if tag == f"task_categories:{task}" or tag == f"task_ids:{task}"
        ]
        fuzzy = matches_term(blob, task)
        if exact:
            score += 5
            reasons.append(f"task:{task}")
        elif fuzzy:
            score += 2
            reasons.append(f"task~:{task}")

    lower_modality_tags = [m.lower() for m in modality_tags]
    for modality in modalities:
        modality = modality.lower().strip()
        if not modality:
            continue
        if f"modality:{modality}" in lower_modality_tags:
            score += 4
            reasons.append(f"modality:{modality}")

    return {
        "dataset_id": dataset.get("id"),
        "score": score,
        "reasons": reasons,
        "task_tags": task_tags,
        "modality_tags": modality_tags,
        "benchmark_tags": collect_prefixed_tags(dataset, ["benchmark:"]),
        "pretty_name": (dataset.get("cardData") or {}).get("pretty_name"),
        "downloads": dataset.get("downloads"),
        "description": " ".join((dataset.get("description") or "").split()),
    }


def search_benchmarks(
    queries: list[str],
    aliases: list[str],
    tasks: list[str],
    modalities: list[str],
    limit: int,
) -> list[dict[str, Any]]:
    datasets = benchmark_catalog(limit=500)
    alias_map = expand_aliases(aliases)

    results = [score_dataset(ds, queries, alias_map, tasks, modalities) for ds in datasets]

    active_filters = bool(queries or aliases or tasks or modalities)
    if active_filters:
        results = [row for row in results if row["score"] >= 2]

    results.sort(
        key=lambda row: (
            -row["score"],
            -(row["downloads"] or 0),
            row["dataset_id"] or "",
        )
    )
    return results[:limit]


def parse_repo_id(repo_id: str) -> tuple[str, str]:
    if "/" not in repo_id:
        raise ValueError(f"Expected <namespace>/<repo>, got: {repo_id}")
    namespace, repo = repo_id.split("/", 1)
    return namespace, repo


def get_leaderboard(repo_id: str, task_id: str | None = None) -> list[dict[str, Any]]:
    namespace, repo = parse_repo_id(repo_id)
    params = {"task_id": task_id} if task_id else None
    data = http_get_json(f"/api/datasets/{namespace}/{repo}/leaderboard", params=params)
    if not isinstance(data, list):
        raise HfApiError(f"Unexpected leaderboard response for {repo_id}")

    normalized: list[dict[str, Any]] = []
    for row in data:
        source = row.get("source") or {}
        normalized.append(
            {
                "dataset_id": repo_id,
                "task_id": task_id,
                "rank": row.get("rank"),
                "model_id": row.get("modelId"),
                "value": row.get("value"),
                "verified": row.get("verified"),
                "lower_is_better": row.get("lower_is_better"),
                "filename": row.get("filename"),
                "notes": row.get("notes"),
                "pull_request": row.get("pullRequest"),
                "source_name": source.get("name"),
                "source_url": source.get("url"),
                "source_is_external": source.get("isExternal"),
            }
        )
    return normalized


def read_repo_ids_from_stdin() -> list[str]:
    if sys.stdin.isatty():
        return []

    repo_ids: list[str] = []
    for raw_line in sys.stdin:
        line = raw_line.strip()
        if not line:
            continue
        if line.startswith("{"):
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue
            candidate = obj.get("dataset_id") or obj.get("id")
            if isinstance(candidate, str) and "/" in candidate:
                repo_ids.append(candidate)
            continue
        if "/" in line:
            repo_ids.append(line)
    return repo_ids


def print_json(data: Any) -> None:
    json.dump(data, sys.stdout, indent=2, ensure_ascii=False)
    sys.stdout.write("\n")


def print_ndjson(rows: list[dict[str, Any]]) -> None:
    for row in rows:
        sys.stdout.write(json.dumps(row, ensure_ascii=False) + "\n")


def print_search_table(rows: list[dict[str, Any]]) -> None:
    if not rows:
        print("No benchmark datasets matched.")
        return

    headers = ["dataset_id", "score", "modalities", "tasks", "reasons", "description"]
    widths = [34, 5, 18, 24, 30, 68]
    print("  ".join(h.ljust(w) for h, w in zip(headers, widths)))
    print("  ".join("-" * w for w in widths))
    for row in rows:
        values = [
            shorten(row.get("dataset_id") or "", widths[0]),
            str(row.get("score", "")),
            shorten(", ".join(row.get("modality_tags") or []), widths[2]),
            shorten(", ".join(row.get("task_tags") or []), widths[3]),
            shorten(", ".join(row.get("reasons") or []), widths[4]),
            shorten(row.get("description") or "", widths[5]),
        ]
        print("  ".join(v.ljust(w) for v, w in zip(values, widths)))


def print_leaderboard_table(rows: list[dict[str, Any]]) -> None:
    if not rows:
        print("No leaderboard rows returned.")
        return

    headers = ["dataset_id", "rank", "model_id", "value", "verified", "source"]
    widths = [30, 5, 38, 10, 8, 28]
    print("  ".join(h.ljust(w) for h, w in zip(headers, widths)))
    print("  ".join("-" * w for w in widths))
    for row in rows:
        values = [
            shorten(str(row.get("dataset_id") or ""), widths[0]),
            str(row.get("rank") or ""),
            shorten(str(row.get("model_id") or ""), widths[2]),
            shorten(str(row.get("value") or ""), widths[3]),
            str(row.get("verified")),
            shorten(str(row.get("source_name") or ""), widths[5]),
        ]
        print("  ".join(v.ljust(w) for v, w in zip(values, widths)))


def build_parser() -> argparse.ArgumentParser:
    parser = FullHelpArgumentParser(
        prog="hf_benchmarks.py",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent(
            """
            Search benchmark datasets and fetch leaderboard results from the Hugging Face Hub.

            Workflow ideas:
              1) Discover candidate benchmarks:
                   hf_benchmarks.py search --alias ocr
                   hf_benchmarks.py search --alias coding
                   hf_benchmarks.py search --task image-to-text --modality document

              2) Inspect a leaderboard:
                   hf_benchmarks.py leaderboard allenai/olmOCR-bench --top 10

              3) Chain search -> leaderboard:
                   hf_benchmarks.py search --alias coding --format ndjson \\
                     | hf_benchmarks.py leaderboard --stdin --top 5 --format table
            """
        ),
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    search_parser = subparsers.add_parser(
        "search",
        help="Search benchmark datasets by query, alias, task, and modality",
    )
    search_parser.add_argument(
        "--query",
        action="append",
        default=[],
        help="Free-text query to match against benchmark dataset metadata. Repeatable.",
    )
    search_parser.add_argument(
        "--alias",
        action="append",
        default=[],
        help=(
            "Convenience alias for common benchmark domains. Known aliases: "
            + ", ".join(sorted(ALIASES))
            + ". Repeatable."
        ),
    )
    search_parser.add_argument(
        "--task",
        action="append",
        default=[],
        help="Task to match, e.g. text-generation, image-to-text, question-answering. Repeatable.",
    )
    search_parser.add_argument(
        "--modality",
        action="append",
        default=[],
        help="Modality to match, e.g. text, image, document, audio. Repeatable.",
    )
    search_parser.add_argument(
        "--limit",
        type=int,
        default=20,
        help="Maximum number of rows to print (default: 20).",
    )
    search_parser.add_argument(
        "--format",
        choices=["table", "json", "ndjson"],
        default="table",
        help="Output format (default: table).",
    )

    leaderboard_parser = subparsers.add_parser(
        "leaderboard",
        help="Fetch normalized leaderboard rows for one or more benchmark datasets",
    )
    leaderboard_parser.add_argument(
        "datasets",
        nargs="*",
        help="Dataset repo ids (<namespace>/<repo>). Can also be supplied via stdin with --stdin.",
    )
    leaderboard_parser.add_argument(
        "--stdin",
        action="store_true",
        help="Read dataset ids from stdin. Accepts plain repo ids or NDJSON with dataset_id/id fields.",
    )
    leaderboard_parser.add_argument(
        "--task-id",
        default=None,
        help="Optional leaderboard task_id query parameter.",
    )
    leaderboard_parser.add_argument(
        "--top",
        type=int,
        default=None,
        help="Only keep the top N results per leaderboard.",
    )
    leaderboard_parser.add_argument(
        "--format",
        choices=["table", "json", "ndjson"],
        default="table",
        help="Output format (default: table).",
    )

    parser._search_parser = search_parser
    parser._leaderboard_parser = leaderboard_parser

    return parser


def run_search(args: argparse.Namespace) -> int:
    rows = search_benchmarks(
        queries=args.query,
        aliases=args.alias,
        tasks=args.task,
        modalities=args.modality,
        limit=args.limit,
    )

    if args.format == "json":
        print_json(rows)
    elif args.format == "ndjson":
        print_ndjson(rows)
    else:
        print_search_table(rows)
    return 0


def run_leaderboard(args: argparse.Namespace) -> int:
    repo_ids = list(args.datasets)
    if args.stdin:
        repo_ids.extend(read_repo_ids_from_stdin())

    deduped: list[str] = []
    seen: set[str] = set()
    for repo_id in repo_ids:
        if repo_id not in seen:
            deduped.append(repo_id)
            seen.add(repo_id)
    repo_ids = deduped

    if not repo_ids:
        print("Error: provide dataset ids or use --stdin.", file=sys.stderr)
        return 2

    rows: list[dict[str, Any]] = []
    for repo_id in repo_ids:
        dataset_rows = get_leaderboard(repo_id, task_id=args.task_id)
        if args.top is not None:
            dataset_rows = dataset_rows[: args.top]
        rows.extend(dataset_rows)

    if args.format == "json":
        print_json(rows)
    elif args.format == "ndjson":
        print_ndjson(rows)
    else:
        print_leaderboard_table(rows)
    return 0


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    try:
        if args.command == "search":
            return run_search(args)
        if args.command == "leaderboard":
            return run_leaderboard(args)
        parser.error(f"Unknown command: {args.command}")
        return 2
    except HfApiError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1
    except ValueError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())