gpu-rent/src/gpu_rent/civitai.py

"""Civitai site API. Bearer only on civitai.com / .red / .green."""

from __future__ import annotations

from dataclasses import dataclass

import httpx

from gpu_rent.errors import CloudError

ALLOWED_HOSTS = ("civitai.com", "civitai.red", "civitai.green")


@dataclass
class CivitaiProbe:
    host: str
    ok: bool
    status: int | None
    detail: str


def _normalize_host(host: str) -> str:
    h = host.strip().lower().removeprefix("https://").removeprefix("http://").split("/")[0]
    if h.startswith("www."):
        h = h[4:]
    return h


def other_host(host: str) -> str:
    h = _normalize_host(host)
    if h.endswith(".red") or h == "civitai.red":
        return "civitai.com"
    return "civitai.red"


def probe_me(token: str, host: str, timeout: float = 15.0) -> CivitaiProbe:
    host = _normalize_host(host)
    if host not in ALLOWED_HOSTS:
        return CivitaiProbe(host=host, ok=False, status=None, detail="хост не из allow-list")
    url = f"https://{host}/api/v1/me"
    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 CivitaiProbe(host=host, ok=False, status=None, detail=str(exc))
    if response.status_code == 200:
        return CivitaiProbe(host=host, ok=True, status=200, detail="токен принят")
    if response.status_code in {401, 403}:
        return CivitaiProbe(
            host=host,
            ok=False,
            status=response.status_code,
            detail="токен отвергнут — перевыпусти ключ на civitai.com/user/account",
        )
    return CivitaiProbe(
        host=host,
        ok=False,
        status=response.status_code,
        detail=response.text[:200] or response.reason_phrase,
    )


def pick_primary_file(version: dict) -> dict | None:
    files = version.get("files") or []
    if not isinstance(files, list):
        return None
    for item in files:
        if isinstance(item, dict) and item.get("primary"):
            return item
    for item in files:
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or "").lower()
        if "safetensor" in name or name.endswith(".safetensors"):
            return item
    for item in files:
        if isinstance(item, dict):
            return item
    return None


def fetch_model_version(token: str, host: str, version_id: int, timeout: float = 30.0) -> tuple[str, dict]:
    """GET /api/v1/model-versions/{id}; one retry on the other Civitai host."""
    first = _normalize_host(host)
    order = [first, other_host(first)]
    last_error = "нет ответа"
    seen: set[str] = set()
    for candidate in order:
        if candidate in seen or candidate not in ALLOWED_HOSTS:
            continue
        seen.add(candidate)
        url = f"https://{candidate}/api/v1/model-versions/{version_id}"
        try:
            with httpx.Client(timeout=timeout, follow_redirects=True) as client:
                headers = {"Authorization": f"Bearer {token}"} if token else {}
                response = client.get(url, headers=headers)
        except httpx.HTTPError as exc:
            last_error = str(exc)
            continue
        if response.status_code == 200:
            data = response.json()
            if isinstance(data, dict) and pick_primary_file(data):
                return candidate, data
            last_error = "пустой files[]"
            continue
        last_error = f"HTTP {response.status_code}"
        if response.status_code not in {404, 400}:
            break
    raise CloudError(f"Civitai version {version_id}: {last_error} (хосты {', '.join(seen)})")


def fetch_model_version_by_hash(
    token: str | None,
    host: str,
    sha256: str,
    timeout: float = 30.0,
) -> tuple[str, dict]:
    """GET /api/v1/model-versions/by-hash/{sha}; public, token optional for NSFW/region."""
    digest = sha256.strip().lower()
    if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest):
        raise CloudError(f"Civitai by-hash: неверный SHA256 ({sha256[:16]}…)")
    first = _normalize_host(host)
    order = [first, other_host(first)]
    last_error = "нет ответа"
    seen: set[str] = set()
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    for candidate in order:
        if candidate in seen or candidate not in ALLOWED_HOSTS:
            continue
        seen.add(candidate)
        url = f"https://{candidate}/api/v1/model-versions/by-hash/{digest}"
        try:
            with httpx.Client(timeout=timeout, follow_redirects=True) as client:
                response = client.get(url, headers=headers)
        except httpx.HTTPError as exc:
            last_error = str(exc)
            continue
        if response.status_code == 200:
            data = response.json()
            if isinstance(data, dict) and data.get("id") is not None:
                return candidate, data
            last_error = "пустой ответ"
            continue
        last_error = f"HTTP {response.status_code}"
        if response.status_code not in {404, 400}:
            break
    raise CloudError(f"Civitai by-hash {digest[:12]}…: {last_error} (хосты {', '.join(seen)})")


def civitai_model_url(model_id: int, version_id: int, host: str = "civitai.red") -> str:
    """Canonical manifest URL (links only — no download)."""
    h = _normalize_host(host)
    if h not in ALLOWED_HOSTS:
        h = "civitai.red"
    if h == "civitai.green":
        h = "civitai.com"
    return f"https://{h}/models/{int(model_id)}?modelVersionId={int(version_id)}"


def version_ids_from_payload(version: dict) -> tuple[int | None, int | None]:
    """Extract (version_id, model_id) from a Civitai version JSON object."""
    try:
        vid = int(version["id"]) if version.get("id") is not None else None
    except (TypeError, ValueError, KeyError):
        vid = None
    mid = version.get("modelId")
    if mid is None and isinstance(version.get("model"), dict):
        mid = version["model"].get("id")
    try:
        model_id = int(mid) if mid is not None else None
    except (TypeError, ValueError):
        model_id = None
    return vid, model_id