src/affix_huggingface/trial.py

"""Trial install registration and the local credential cache.

The proof allowance is counted on api.affix-io.com against a subject digest.
Files written here are a cache. Deleting them forces re-registration against
the same server-side counter and does not restore spent proofs.
"""

# Credit: @paparichens

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .errors import ConfigurationError

TRIAL_PRODUCT = "huggingface"

# Published trial key. It ships inside the package, so treat it as public:
# it is scoped, rate limited, and revocable. Supply AFFIX_API_KEY for
# production use.
PUBLIC_TRIAL_API_KEY = "aio_50754e6f66c7a573a46a16a52e1b1edffba48c395b708316"


@dataclass(frozen=True)
class TrialCredentials:
    """Install credentials issued by AffixIO and cached locally."""

    install_id: str
    install_secret: str
    subject_id: str
    product: str = TRIAL_PRODUCT

    def headers(self) -> dict[str, str]:
        return {
            "X-Affix-Install-Id": self.install_id,
            "X-Affix-Install-Secret": self.install_secret,
        }

    def to_dict(self) -> dict[str, str]:
        return {
            "install_id": self.install_id,
            "install_secret": self.install_secret,
            "subject_id": self.subject_id,
            "product": self.product,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> TrialCredentials | None:
        try:
            return cls(
                install_id=str(data["install_id"]),
                install_secret=str(data["install_secret"]),
                subject_id=str(data["subject_id"]),
                product=str(data.get("product", TRIAL_PRODUCT)),
            )
        except (KeyError, TypeError):
            return None


@dataclass(frozen=True)
class TrialQuota:
    """Allowance reported by AffixIO."""

    limit: int
    used: int
    remaining: int

    @property
    def exhausted(self) -> bool:
        return self.remaining <= 0

    @classmethod
    def from_headers(cls, headers: Any) -> TrialQuota | None:
        try:
            limit = headers.get("X-Affix-Trial-Limit")
            used = headers.get("X-Affix-Trial-Used")
            remaining = headers.get("X-Affix-Trial-Remaining")
        except AttributeError:
            return None
        if limit is None or used is None or remaining is None:
            return None
        try:
            return cls(limit=int(limit), used=int(used), remaining=int(remaining))
        except (TypeError, ValueError):
            return None


def credentials_path() -> Path:
    """Cache location for install credentials."""
    override = os.getenv("AFFIX_INSTALL_FILE")
    if override:
        return Path(override).expanduser()
    base = os.getenv("XDG_CACHE_HOME")
    root = Path(base).expanduser() if base else Path.home() / ".cache"
    return root / "affix-huggingface" / "install.json"


def load_credentials(path: Path | None = None) -> TrialCredentials | None:
    target = path or credentials_path()
    try:
        raw = json.loads(target.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    if not isinstance(raw, dict):
        return None
    return TrialCredentials.from_dict(raw)


def save_credentials(credentials: TrialCredentials, path: Path | None = None) -> None:
    target = path or credentials_path()
    try:
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(
            json.dumps(credentials.to_dict(), indent=2) + "\n",
            encoding="utf-8",
        )
        target.chmod(0o600)
    except OSError:
        # A read-only cache directory costs a registration round trip per run,
        # it does not grant extra proofs.
        return


def subject_digest(raw: str) -> str:
    """Hash a subject locally so the identity itself never leaves the host."""
    value = raw.strip().lower()
    if not value:
        raise ConfigurationError("Trial subject must not be empty")
    return hashlib.sha256(f"affix-hf-subject:{value}".encode()).hexdigest()


def resolve_subject(
    *,
    subject: str | None = None,
    hf_token: str | None = None,
) -> str:
    """
    Determine the identity the allowance is counted against.

    Order of preference: explicit subject, AFFIX_TRIAL_SUBJECT, then the
    Hugging Face account behind the token. Resolution fails closed, because a
    random fallback would hand out a fresh allowance on every install.
    """
    explicit = subject or os.getenv("AFFIX_TRIAL_SUBJECT")
    if explicit:
        return subject_digest(explicit)

    token = hf_token or os.getenv("HF_TOKEN")
    if not token:
        raise ConfigurationError(
            "Cannot determine a trial subject. Set HF_TOKEN so the Hugging Face "
            "account can be resolved, or pass subject= explicitly."
        )

    try:
        from huggingface_hub import whoami

        info = whoami(token=token)
    except Exception as exc:
        raise ConfigurationError(
            "Could not resolve the Hugging Face account for the trial allowance. "
            "Pass subject= explicitly if this host cannot reach huggingface.co."
        ) from exc

    name = None
    if isinstance(info, dict):
        name = info.get("id") or info.get("name")
    else:
        name = getattr(info, "id", None) or getattr(info, "name", None)
    if not name:
        raise ConfigurationError("Hugging Face account lookup returned no identity")
    return subject_digest(f"hf:{name}")