beta_keys_store.py

"""
VTX beta keys — private Hugging Face dataset storage with local cache.

Environment:
  VTX_KEYS_SOURCE=local|dataset          (default: dataset)
  VTX_KEYS_PATH=path/to/keys.json        (local mode / cache file override)
  VTX_BETA_KEYS_DATASET=org/repo         (default: ManChildTechnologies/VTX-BetaKeys)
  VTX_BETA_KEYS_FILENAME=keys.json
  VTX_BETA_KEYS_REVISION=main
  VTX_KEYS_REFRESH_SECONDS=30
  HF_TOKEN / HUGGING_FACE_HUB_TOKEN
"""
from __future__ import annotations

import io
import json
import os
import threading
import time
from pathlib import Path
from typing import Any

from beta_keys_schema import migrate_keys_payload, normalize_record

APP_DIR = Path(__file__).resolve().parent
DEFAULT_KEYS_PATH = APP_DIR / "keys.json"
DEFAULT_CACHE_PATH = APP_DIR / "keys_cache.json"
DEFAULT_DATASET_REPO = "ManChildTechnologies/VTX-BetaKeys"
DEFAULT_DATASET_FILENAME = "keys.json"

_store_lock = threading.RLock()
_keys_cache: dict[str, Any] = {
    "source": None,
    "signature": None,
    "fetched_at": 0.0,
    "data": {},
    "error": "",
    "using_stale_cache": False,
}


def keys_source() -> str:
    raw = str(os.environ.get("VTX_KEYS_SOURCE", "dataset") or "dataset").strip().lower()
    return "local" if raw == "local" else "dataset"


def keys_path() -> Path:
    custom = str(os.environ.get("VTX_KEYS_PATH", "") or "").strip()
    if custom:
        return Path(custom)
    return DEFAULT_KEYS_PATH


def cache_path() -> Path:
    custom = str(os.environ.get("VTX_KEYS_CACHE_PATH", "") or "").strip()
    if custom:
        return Path(custom)
    return DEFAULT_CACHE_PATH


def dataset_repo() -> str:
    return str(os.environ.get("VTX_BETA_KEYS_DATASET", DEFAULT_DATASET_REPO) or DEFAULT_DATASET_REPO).strip()


def dataset_filename() -> str:
    return str(os.environ.get("VTX_BETA_KEYS_FILENAME", DEFAULT_DATASET_FILENAME) or DEFAULT_DATASET_FILENAME).strip()


def dataset_revision() -> str:
    return str(os.environ.get("VTX_BETA_KEYS_REVISION", "main") or "main").strip() or "main"


def refresh_seconds() -> int:
    try:
        return max(5, int(os.environ.get("VTX_KEYS_REFRESH_SECONDS", "30") or 30))
    except Exception:
        return 30


def hf_token() -> str:
    for key in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN"):
        value = str(os.environ.get(key, "") or "").strip()
        if value:
            return value
    return ""


def _require_hf_token() -> str:
    token = hf_token()
    if not token:
        raise RuntimeError(
            "HF_TOKEN is required for VTX beta key dataset access. "
            "Set HF_TOKEN as an environment variable or Hugging Face Space secret."
        )
    return token


def _get_hf_api():
    try:
        from huggingface_hub import HfApi
    except ImportError as exc:
        raise RuntimeError("huggingface_hub is required for dataset key storage.") from exc
    return HfApi(token=_require_hf_token())


def _serialize_keys(keys: dict[str, dict[str, Any]]) -> str:
    ordered = {
        key: normalize_record(key, record)
        for key, record in sorted(keys.items(), key=lambda item: item[0])
    }
    return json.dumps(ordered, indent=2) + "\n"


def _read_json_file(path: Path) -> dict[str, dict[str, Any]]:
    if not path.is_file():
        return {}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    return migrate_keys_payload(payload)


def _write_json_file(path: Path, keys: dict[str, dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(_serialize_keys(keys), encoding="utf-8")


def _local_signature(path: Path) -> str | None:
    try:
        if not path.is_file():
            return None
        stat = path.stat()
        return f"local:{path}:{stat.st_mtime_ns}:{stat.st_size}"
    except OSError:
        return None


def _fetch_dataset_keys() -> tuple[dict[str, dict[str, Any]], str]:
    token = _require_hf_token()
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise RuntimeError("huggingface_hub is required for dataset key loading.") from exc

    repo_id = dataset_repo()
    filename = dataset_filename()
    revision = dataset_revision()
    cached_path = hf_hub_download(
        repo_id=repo_id,
        repo_type="dataset",
        filename=filename,
        revision=revision,
        token=token,
        force_download=True,
    )
    path = Path(cached_path)
    data = _read_json_file(path)
    signature = f"dataset:{repo_id}:{revision}:{filename}:{path.stat().st_mtime_ns}:{path.stat().st_size}"
    _write_json_file(cache_path(), data)
    return data, signature


def push_keys_to_dataset(keys: dict[str, dict[str, Any]], commit_message: str) -> None:
    token = _require_hf_token()
    api = _get_hf_api()
    payload = _serialize_keys(keys).encode("utf-8")
    api.upload_file(
        path_or_fileobj=io.BytesIO(payload),
        path_in_repo=dataset_filename(),
        repo_id=dataset_repo(),
        repo_type="dataset",
        revision=dataset_revision(),
        commit_message=commit_message,
        token=token,
    )
    _write_json_file(cache_path(), keys)
    global _keys_cache
    _keys_cache = {
        "source": "dataset",
        "signature": f"pushed:{time.time_ns()}",
        "fetched_at": time.time(),
        "data": dict(keys),
        "error": "",
        "using_stale_cache": False,
    }


def load_keys(force: bool = False) -> dict[str, dict[str, Any]]:
    global _keys_cache

    with _store_lock:
        source = keys_source()
        now = time.time()
        ttl = refresh_seconds()

        if source == "local":
            path = keys_path()
            signature = _local_signature(path)
            if (
                not force
                and _keys_cache.get("source") == "local"
                and signature == _keys_cache.get("signature")
                and (now - float(_keys_cache.get("fetched_at") or 0.0)) < ttl
            ):
                return dict(_keys_cache.get("data") or {})

            data = _read_json_file(path)
            _keys_cache = {
                "source": "local",
                "signature": signature,
                "fetched_at": now,
                "data": data,
                "error": "",
                "using_stale_cache": False,
            }
            return dict(data)

        if (
            not force
            and _keys_cache.get("source") == "dataset"
            and (now - float(_keys_cache.get("fetched_at") or 0.0)) < ttl
            and _keys_cache.get("data") is not None
        ):
            return dict(_keys_cache.get("data") or {})

        try:
            data, signature = _fetch_dataset_keys()
            _keys_cache = {
                "source": "dataset",
                "signature": signature,
                "fetched_at": now,
                "data": data,
                "error": "",
                "using_stale_cache": False,
            }
            return dict(data)
        except Exception as exc:
            cached_file = _read_json_file(cache_path())
            if cached_file:
                _keys_cache = {
                    "source": "dataset",
                    "signature": _keys_cache.get("signature"),
                    "fetched_at": now,
                    "data": cached_file,
                    "error": str(exc),
                    "using_stale_cache": True,
                }
                return dict(cached_file)
            _keys_cache = {
                "source": "dataset",
                "signature": None,
                "fetched_at": now,
                "data": {},
                "error": str(exc),
                "using_stale_cache": False,
            }
            raise


def save_keys(keys: dict[str, dict[str, Any]], commit_message: str = "Update VTX beta keys") -> None:
    normalized = {
        key: normalize_record(key, record)
        for key, record in (keys or {}).items()
        if str(key or "").strip()
    }
    with _store_lock:
        if keys_source() == "local":
            _write_json_file(keys_path(), normalized)
            global _keys_cache
            _keys_cache = {
                "source": "local",
                "signature": _local_signature(keys_path()),
                "fetched_at": time.time(),
                "data": dict(normalized),
                "error": "",
                "using_stale_cache": False,
            }
            return
        push_keys_to_dataset(normalized, commit_message)


def update_key_record(
    beta_key: str,
    updater,
    commit_message: str,
) -> dict[str, dict[str, Any]]:
    key = str(beta_key or "").strip()
    if not key:
        raise ValueError("Beta key is required.")

    with _store_lock:
        keys = load_keys(force=True)
        record = dict(keys.get(key) or normalize_record(key, {"key": key}))
        record = normalize_record(key, updater(record))
        record["key"] = key
        keys[key] = record
        save_keys(keys, commit_message=commit_message)
        return record


def keys_status() -> dict[str, Any]:
    source = keys_source()
    status = {
        "source": source,
        "keys_loaded": len(_keys_cache.get("data") or {}),
        "last_error": str(_keys_cache.get("error") or ""),
        "cache_age_seconds": max(0, int(time.time() - float(_keys_cache.get("fetched_at") or 0.0))),
        "using_stale_cache": bool(_keys_cache.get("using_stale_cache")),
    }
    if source == "local":
        path = keys_path()
        status["keys_path"] = str(path)
        status["keys_exists"] = path.is_file()
    else:
        status["dataset_repo"] = dataset_repo()
        status["dataset_filename"] = dataset_filename()
        status["dataset_revision"] = dataset_revision()
        status["hf_token_configured"] = bool(hf_token())
        status["cache_path"] = str(cache_path())
    return status