backend/search.py

"""Hugging Face dataset search and evidence collection."""
from __future__ import annotations

import json
import os
from typing import Any
from urllib.parse import urlencode
from urllib.request import Request, urlopen

from huggingface_hub import HfApi

_api = None
_VIEWER_BASE = "https://datasets-server.huggingface.co"


def _get_api() -> HfApi:
    global _api
    if _api is None:
        token = os.getenv("HF_TOKEN")
        _api = HfApi(token=token if token else None)
    return _api


def _tag_values(tags: list[str], prefix: str) -> list[str]:
    marker = f"{prefix}:"
    return [tag[len(marker):] for tag in tags if tag.startswith(marker)]


def _http_json(path: str, params: dict[str, str], timeout: float = 8.0) -> dict[str, Any]:
    url = f"{_VIEWER_BASE}{path}?{urlencode(params)}"
    request = Request(url, headers={"User-Agent": "hf-agentic-search/1.0"})
    with urlopen(request, timeout=timeout) as response:
        return json.load(response)


def _normalize_dataset(ds: Any) -> dict[str, Any]:
    tags = list(getattr(ds, "tags", []) or [])
    dataset_id = ds.id
    return {
        "id": dataset_id,
        "author": (getattr(ds, "author", "") or "")
        or (dataset_id.split("/")[0] if "/" in dataset_id else ""),
        "tags": tags,
        "downloads": getattr(ds, "downloads", 0) or 0,
        "likes": getattr(ds, "likes", 0) or 0,
        "task_categories": list(getattr(ds, "task_categories", []) or [])
        or _tag_values(tags, "task_categories"),
        "languages": [str(value) for value in (getattr(ds, "languages", []) or [])]
        or _tag_values(tags, "language"),
        "license": str(getattr(ds, "license", "") or "")
        or (_tag_values(tags, "license") or [""])[0],
        "size_category": str(getattr(ds, "size_category", "") or "")
        or (_tag_values(tags, "size_categories") or [""])[0],
        "formats": _tag_values(tags, "format"),
        "modalities": _tag_values(tags, "modality"),
        "description": str(getattr(ds, "description", "") or "")[:1200],
        "created_at": str(getattr(ds, "created_at", "") or ""),
        "updated_at": str(getattr(ds, "updated_at", "") or ""),
        "private": bool(getattr(ds, "private", False)),
        "gated": getattr(ds, "gated", False) or False,
        "hub_url": f"https://huggingface.co/datasets/{dataset_id}",
    }


def search_datasets(query: str, limit: int = 20) -> list[dict[str, Any]]:
    """Search the Hub and normalize metadata encoded in dataset tags."""
    try:
        results = list(_get_api().list_datasets(search=query, limit=limit))
    except Exception:
        return []
    datasets = []
    for dataset in results:
        try:
            datasets.append(_normalize_dataset(dataset))
        except Exception:
            continue
    return datasets


def inspect_dataset(dataset_id: str, base: dict[str, Any] | None = None) -> dict[str, Any]:
    """Collect card, config, split, feature, and sample evidence for one dataset."""
    evidence = dict(base or {"id": dataset_id})
    evidence.update(
        {
            "accessible": False,
            "inspection_error": "",
            "configs": [],
            "splits": [],
            "features": [],
            "sample_rows": [],
            "files": [],
            "card_complete": False,
            "num_examples": 0,
        }
    )
    try:
        info = _get_api().dataset_info(dataset_id, files_metadata=True)
        evidence.update(_normalize_dataset(info))
        evidence["accessible"] = not bool(getattr(info, "private", False))
        evidence["files"] = [
            sibling.rfilename
            for sibling in (getattr(info, "siblings", []) or [])
            if getattr(sibling, "rfilename", None)
        ][:40]
        card_data = getattr(info, "card_data", None)
        if card_data:
            card_dict = card_data.to_dict() if hasattr(card_data, "to_dict") else dict(card_data)
            evidence["card_data"] = card_dict
            evidence["card_complete"] = bool(
                evidence.get("description")
                and (card_dict.get("license") or evidence.get("license"))
                and (card_dict.get("language") or evidence.get("languages"))
            )
    except Exception as exc:
        evidence["inspection_error"] = f"Hub metadata unavailable: {type(exc).__name__}"
        return evidence

    try:
        viewer = _http_json("/info", {"dataset": dataset_id})
        dataset_info = viewer.get("dataset_info", {})
        for config_name, config_data in dataset_info.items():
            evidence["configs"].append(config_name)
            splits = config_data.get("splits", {}) or {}
            split_items = splits.values() if isinstance(splits, dict) else splits
            for split in split_items:
                split_name = split.get("name") if isinstance(split, dict) else None
                if split_name and split_name not in evidence["splits"]:
                    evidence["splits"].append(split_name)
                if isinstance(split, dict):
                    evidence["num_examples"] += int(split.get("num_examples") or 0)
            features = config_data.get("features") or {}
            if isinstance(features, dict):
                feature_names = list(features)
            elif isinstance(features, list):
                feature_names = [
                    item.get("name") for item in features if isinstance(item, dict) and item.get("name")
                ]
            else:
                feature_names = []
            for feature in feature_names:
                if feature not in evidence["features"]:
                    evidence["features"].append(feature)
    except Exception:
        pass

    if evidence["configs"] and evidence["splits"]:
        try:
            rows = _http_json(
                "/first-rows",
                {
                    "dataset": dataset_id,
                    "config": evidence["configs"][0],
                    "split": evidence["splits"][0],
                },
                timeout=10.0,
            )
            evidence["sample_rows"] = [
                item.get("row", {}) for item in (rows.get("rows", []) or [])[:3]
            ]
            row_features = rows.get("features") or []
            for feature in row_features:
                name = feature.get("name") if isinstance(feature, dict) else None
                if name and name not in evidence["features"]:
                    evidence["features"].append(name)
            if not evidence["features"] and evidence["sample_rows"]:
                evidence["features"] = list(evidence["sample_rows"][0])
        except Exception:
            pass

    evidence["configs"] = evidence["configs"][:10]
    evidence["splits"] = evidence["splits"][:12]
    evidence["features"] = evidence["features"][:30]
    return evidence


def get_dataset_info(dataset_id: str) -> dict[str, Any] | None:
    """Backward-compatible detailed dataset lookup."""
    result = inspect_dataset(dataset_id)
    return result if result.get("accessible") else None