gpu-rent/src/gpu_rent/huggingface.py

"""Hugging Face Hub: auth probe, metadata lookup (capture fallback), URL helpers."""

from __future__ import annotations

from dataclasses import dataclass
from urllib.parse import quote, urlparse

import httpx

from gpu_rent.errors import CloudError

HF_API = "https://huggingface.co/api"
HF_HOST = "huggingface.co"


@dataclass
class HfProbe:
    ok: bool
    name: str = ""
    detail: str = ""


@dataclass
class HfFileHit:
    repo_id: str
    filename: str
    url: str
    sha256: str = ""
    title: str = ""


def is_huggingface_url(url: str) -> bool:
    host = (urlparse(url).hostname or "").lower()
    return host == HF_HOST or host.endswith(".huggingface.co")


def hf_resolve_url(repo_id: str, filename: str, *, revision: str = "main") -> str:
    repo = repo_id.strip().strip("/")
    name = filename.lstrip("/")
    rev = revision.strip() or "main"
    return f"https://{HF_HOST}/{repo}/resolve/{quote(rev, safe='')}/{quote(name, safe='/')}"


def probe_whoami(token: str, timeout: float = 15.0) -> HfProbe:
    token = (token or "").strip()
    if not token:
        return HfProbe(ok=False, detail="токена нет")
    url = f"{HF_API}/whoami-v2"
    try:
        with httpx.Client(timeout=timeout, follow_redirects=True) as client:
            response = client.get(url, headers={"Authorization": f"Bearer {token}"})
    except httpx.HTTPError as exc:
        return HfProbe(ok=False, detail=str(exc)[:160])
    if response.status_code == 200:
        data = response.json() if response.content else {}
        name = ""
        if isinstance(data, dict):
            name = str(data.get("name") or data.get("fullname") or "")
        return HfProbe(ok=True, name=name, detail="токен принят")
    if response.status_code in {401, 403}:
        return HfProbe(
            ok=False,
            detail="токен отвергнут — https://huggingface.co/settings/tokens",
        )
    return HfProbe(
        ok=False,
        detail=f"HTTP {response.status_code}: {(response.text or '')[:120]}",
    )


def _headers(token: str | None) -> dict[str, str]:
    h = {"User-Agent": "gpu-rent/1"}
    if token:
        h["Authorization"] = f"Bearer {token}"
    return h


def _sha_from_entry(entry: dict) -> str:
    lfs = entry.get("lfs")
    if isinstance(lfs, dict):
        for key in ("sha256", "oid", "sha"):
            val = lfs.get(key)
            if val:
                return str(val).strip().lower()
    for key in ("sha256", "oid", "blob_id"):
        val = entry.get(key)
        if val and len(str(val)) >= 40:
            return str(val).strip().lower()
    return ""


def _tree_match(
    client: httpx.Client,
    repo_id: str,
    want_sha: str,
    *,
    filename_hint: str | None,
    token: str | None,
) -> HfFileHit | None:
    url = f"{HF_API}/models/{repo_id}/tree/main"
    response = client.get(
        url,
        params={"recursive": "true"},
        headers=_headers(token),
    )
    if response.status_code != 200:
        return None
    data = response.json()
    if not isinstance(data, list):
        return None
    hint = (filename_hint or "").replace("\\", "/").split("/")[-1].lower()
    for entry in data:
        if not isinstance(entry, dict):
            continue
        if entry.get("type") and entry.get("type") != "file":
            continue
        path = str(entry.get("path") or entry.get("rfilename") or "")
        if not path:
            continue
        sha = _sha_from_entry(entry)
        if sha and sha == want_sha:
            return HfFileHit(
                repo_id=repo_id,
                filename=path,
                url=hf_resolve_url(repo_id, path),
                sha256=sha,
                title=repo_id,
            )
        # Filename fallback when tree has no sha (rare) — only exact name match
        if hint and path.replace("\\", "/").split("/")[-1].lower() == hint and not sha:
            return HfFileHit(
                repo_id=repo_id,
                filename=path,
                url=hf_resolve_url(repo_id, path),
                sha256="",
                title=repo_id,
            )
    return None


def lookup_by_sha256(
    token: str | None,
    sha256: str,
    *,
    filename: str | None = None,
    limit_repos: int = 12,
    timeout: float = 45.0,
) -> HfFileHit | None:
    """Best-effort: search Hub by filename, match LFS sha256 in repo tree.

    Hugging Face has no public global by-hash index; this is the practical fallback
    after Civitai miss. Token optional for public repos; gated need HF_TOKEN.
    """
    digest = sha256.strip().lower()
    if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest):
        raise CloudError(f"HF by-hash: неверный SHA256 ({sha256[:16]}…)")

    hint = (filename or "").replace("\\", "/").split("/")[-1]
    stem = hint.rsplit(".", 1)[0] if hint else ""
    queries: list[str] = []
    if hint:
        queries.append(hint)
    if stem and stem.lower() != hint.lower():
        queries.append(stem)
    if not queries:
        return None

    try:
        with httpx.Client(timeout=timeout, follow_redirects=True) as client:
            seen: set[str] = set()
            for q in queries:
                response = client.get(
                    f"{HF_API}/models",
                    params={"search": q, "limit": limit_repos},
                    headers=_headers(token),
                )
                if response.status_code in {401, 403}:
                    raise CloudError(
                        "HF search: 401/403 — нужен валидный HF_TOKEN "
                        "(https://huggingface.co/settings/tokens)"
                    )
                if response.status_code != 200:
                    continue
                models = response.json()
                if not isinstance(models, list):
                    continue
                for model in models:
                    if not isinstance(model, dict):
                        continue
                    repo_id = str(model.get("modelId") or model.get("id") or "").strip()
                    if not repo_id or repo_id in seen:
                        continue
                    seen.add(repo_id)
                    hit = _tree_match(
                        client, repo_id, digest, filename_hint=hint, token=token
                    )
                    if hit:
                        return hit
                    if len(seen) >= limit_repos:
                        return None
    except httpx.HTTPError as exc:
        raise CloudError(f"HF lookup: {exc}") from exc
    return None