kata_store.py

from __future__ import annotations

import json
import os
import shutil
import tempfile
import time
from pathlib import Path
from typing import Any, Protocol

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError


META_FIELDS = (
    "id",
    "prompt",
    "seconds",
    "fps",
    "num_frames",
    "model",
    "seed",
    "denoising_steps",
    "created_at",
    "continues_from",
    "heading_offset",
    "created_by",
    "created_by_name",
    "created_by_avatar",
    # MY KARATE composer: the editable recipe (stance ids + transition prompts/dials)
    # and the owning browser, stored on a kata's ROOT record so the library can list
    # + reopen it. Small JSON, kept in the manifest so the list view needs no per-record fetch.
    "compose_recipe",
    "compose_owner",
)


def split_meta(record: dict[str, Any]) -> dict[str, Any]:
    return {k: record[k] for k in META_FIELDS if k in record}


class AnimationStore(Protocol):
    def save(self, record: dict[str, Any], preview: dict[str, Any] | None = None, npz_path: str | None = None) -> str: ...
    def list(self, limit: int = 200) -> list[dict[str, Any]]: ...
    def get(self, record_id: str) -> dict[str, Any] | None: ...
    def get_meta(self, record_id: str) -> dict[str, Any] | None: ...
    def get_preview(self, record_id: str) -> dict[str, Any] | None: ...
    def get_npz(self, record_id: str) -> str | None: ...
    def exists(self, record_id: str) -> bool: ...
    def delete(self, record_id: str) -> bool: ...


class LocalAnimationStore:
    def __init__(self, root: str | Path | None = None) -> None:
        self.root = Path(root or os.environ.get("KIMODO_STORE_PATH", ".kimodo-animations"))
        self.records = self.root / "animations"
        self.previews = self.root / "previews"
        self.npzs = self.root / "npz"
        for path in (self.records, self.previews, self.npzs):
            path.mkdir(parents=True, exist_ok=True)

    def _safe(self, record_id: str) -> str:
        return "".join(c if c.isalnum() or c in "-_" else "_" for c in record_id)

    def _record_path(self, record_id: str) -> Path:
        return self.records / f"{self._safe(record_id)}.json"

    def _preview_path(self, record_id: str) -> Path:
        return self.previews / f"{self._safe(record_id)}.json"

    def _npz_path(self, record_id: str) -> Path:
        return self.npzs / f"{self._safe(record_id)}.npz"

    def save(self, record: dict[str, Any], preview: dict[str, Any] | None = None, npz_path: str | None = None) -> str:
        if "id" not in record:
            raise ValueError("record id is required")
        if "created_at" not in record:
            record["created_at"] = int(time.time())
        record_id = str(record["id"])
        self._record_path(record_id).write_text(json.dumps(record), encoding="utf-8")
        if preview is not None:
            self._preview_path(record_id).write_text(json.dumps(preview), encoding="utf-8")
        if npz_path:
            shutil.copyfile(npz_path, self._npz_path(record_id))
        return record_id

    def list(self, limit: int = 200) -> list[dict[str, Any]]:
        out = []
        files = sorted(self.records.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True)
        for path in files[: max(0, int(limit))]:
            try:
                out.append(split_meta(json.loads(path.read_text(encoding="utf-8"))))
            except Exception:
                continue
        return out

    def get(self, record_id: str) -> dict[str, Any] | None:
        path = self._record_path(record_id)
        if not path.exists():
            return None
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            return None

    def get_meta(self, record_id: str) -> dict[str, Any] | None:
        record = self.get(record_id)
        return split_meta(record) if record else None

    def get_preview(self, record_id: str) -> dict[str, Any] | None:
        path = self._preview_path(record_id)
        if not path.exists():
            record = self.get(record_id)
            return None if record is None else record
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            return None

    def get_npz(self, record_id: str) -> str | None:
        path = self._npz_path(record_id)
        return str(path) if path.exists() else None

    def exists(self, record_id: str) -> bool:
        return self._record_path(record_id).exists()

    def delete(self, record_id: str) -> bool:
        existed = False
        for path in (self._record_path(record_id), self._preview_path(record_id), self._npz_path(record_id)):
            if path.exists():
                path.unlink()
                existed = True
        return existed


class HFDatasetAnimationStore:
    def __init__(self, repo_id: str, token: str | None = None) -> None:
        self.repo_id = repo_id
        self.token = token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        self.api = HfApi(token=self.token)
        self.api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, private=False)

    def _download_json(self, path_in_repo: str) -> dict[str, Any] | None:
        try:
            path = hf_hub_download(
                repo_id=self.repo_id,
                repo_type="dataset",
                filename=path_in_repo,
                token=self.token,
            )
        except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
            return None
        try:
            return json.loads(Path(path).read_text(encoding="utf-8"))
        except Exception:
            return None

    def _record_path(self, record_id: str) -> str:
        return f"animations/{record_id}.json"

    def _preview_path(self, record_id: str) -> str:
        return f"previews/{record_id}.json"

    def _npz_path(self, record_id: str) -> str:
        return f"npz/{record_id}.npz"

    # A single small metadata file (id -> split_meta) so the dropdown can be
    # populated with ONE download instead of fetching every full record (each
    # carries megabytes of quaternion/joint arrays). Kept in sync on save/delete;
    # list() rebuilds it from animations/*.json if it's ever missing or stale.
    MANIFEST_PATH = "manifest.json"

    def _load_manifest(self) -> dict[str, dict[str, Any]]:
        data = self._download_json(self.MANIFEST_PATH)
        if isinstance(data, dict) and isinstance(data.get("records"), dict):
            return dict(data["records"])
        return {}

    def _manifest_op(self, records: dict[str, dict[str, Any]]) -> CommitOperationAdd:
        blob = json.dumps({"records": records}, separators=(",", ":")).encode("utf-8")
        return CommitOperationAdd(path_in_repo=self.MANIFEST_PATH, path_or_fileobj=blob)

    def _scan_meta(self) -> dict[str, dict[str, Any]]:
        """Slow path: read every full record to (re)build the manifest."""
        try:
            files = self.api.list_repo_files(self.repo_id, repo_type="dataset")
        except RepositoryNotFoundError:
            return {}
        out: dict[str, dict[str, Any]] = {}
        for path in files:
            if not path.startswith("animations/") or not path.endswith(".json"):
                continue
            record = self._download_json(path)
            if record and record.get("id"):
                out[str(record["id"])] = split_meta(record)
        return out

    def save(self, record: dict[str, Any], preview: dict[str, Any] | None = None, npz_path: str | None = None) -> str:
        if "id" not in record:
            raise ValueError("record id is required")
        if "created_at" not in record:
            record["created_at"] = int(time.time())
        record_id = str(record["id"])

        operations = []
        record_bytes = json.dumps(record, separators=(",", ":")).encode("utf-8")
        operations.append(CommitOperationAdd(path_in_repo=self._record_path(record_id), path_or_fileobj=record_bytes))
        if preview is not None:
            preview_bytes = json.dumps(preview, separators=(",", ":")).encode("utf-8")
            operations.append(CommitOperationAdd(path_in_repo=self._preview_path(record_id), path_or_fileobj=preview_bytes))
        if npz_path:
            operations.append(CommitOperationAdd(path_in_repo=self._npz_path(record_id), path_or_fileobj=npz_path))

        # Keep the manifest current in the SAME commit (best-effort; a manifest
        # read failure must never block saving the actual animation).
        try:
            manifest = self._load_manifest()
            manifest[record_id] = split_meta(record)
            operations.append(self._manifest_op(manifest))
        except Exception as exc:  # pragma: no cover
            print(f"manifest update skipped on save: {type(exc).__name__}: {exc}")

        self.api.create_commit(
            repo_id=self.repo_id,
            repo_type="dataset",
            operations=operations,
            commit_message=f"Save animation {record_id}",
        )
        return record_id

    def list(self, limit: int = 200) -> list[dict[str, Any]]:
        manifest = self._load_manifest()
        if not manifest:
            # No manifest yet (legacy dataset): rebuild it once from full records.
            manifest = self._scan_meta()
            if manifest:
                try:
                    self.api.create_commit(
                        repo_id=self.repo_id,
                        repo_type="dataset",
                        operations=[self._manifest_op(manifest)],
                        commit_message="Build animation manifest",
                    )
                except Exception as exc:  # pragma: no cover
                    print(f"manifest build skipped: {type(exc).__name__}: {exc}")
        records = list(manifest.values())
        records.sort(key=lambda r: r.get("created_at", 0), reverse=True)
        return records[: max(0, int(limit))]

    def get(self, record_id: str) -> dict[str, Any] | None:
        return self._download_json(self._record_path(record_id))

    def get_meta(self, record_id: str) -> dict[str, Any] | None:
        manifest = self._load_manifest()
        if record_id in manifest:
            return manifest[record_id]
        record = self.get(record_id)
        return split_meta(record) if record else None

    def get_preview(self, record_id: str) -> dict[str, Any] | None:
        preview = self._download_json(self._preview_path(record_id))
        if preview is not None:
            return preview
        return self.get(record_id)

    def get_npz(self, record_id: str) -> str | None:
        try:
            return hf_hub_download(
                repo_id=self.repo_id,
                repo_type="dataset",
                filename=self._npz_path(record_id),
                token=self.token,
                local_dir=Path(tempfile.gettempdir()) / "kimodo_dataset_npz",
            )
        except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
            return None

    def exists(self, record_id: str) -> bool:
        try:
            hf_hub_download(
                repo_id=self.repo_id,
                repo_type="dataset",
                filename=self._record_path(record_id),
                token=self.token,
            )
            return True
        except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
            return False

    def delete(self, record_id: str) -> bool:
        deleted = False
        for path in (self._record_path(record_id), self._preview_path(record_id), self._npz_path(record_id)):
            try:
                self.api.delete_file(
                    repo_id=self.repo_id,
                    repo_type="dataset",
                    path_in_repo=path,
                    commit_message=f"Delete {path}",
                )
                deleted = True
            except Exception:
                continue
        try:
            manifest = self._load_manifest()
            if record_id in manifest:
                del manifest[record_id]
                self.api.create_commit(
                    repo_id=self.repo_id,
                    repo_type="dataset",
                    operations=[self._manifest_op(manifest)],
                    commit_message=f"Remove {record_id} from manifest",
                )
        except Exception as exc:  # pragma: no cover
            print(f"manifest update skipped on delete: {type(exc).__name__}: {exc}")
        return deleted


def make_store() -> AnimationStore:
    backend = os.environ.get("KIMODO_STORE_BACKEND", "").strip().lower()
    dataset_repo = os.environ.get("KIMODO_DATASET_REPO", "").strip()
    if backend == "dataset" or dataset_repo:
        return HFDatasetAnimationStore(dataset_repo or "polats/kimodo-kata-animations")
    return LocalAnimationStore()


store = make_store()