huggingface_human_training_space/studio_store.py

"""Persistent helpers for true-studio run history, drafts, and artifacts."""

from __future__ import annotations

import json
import tempfile
from datetime import UTC, datetime
from pathlib import Path
from threading import Lock
from typing import Any
from uuid import uuid4

STORE_LOCK = Lock()
STORE_FALLBACK_DIRNAME = "maris-human-training-space-studio"
RUNS_FILENAME = "human-training-runs.json"
DRAFTS_FILENAME = "human-training-drafts.json"
ARTIFACTS_FILENAME = "human-training-artifacts.json"


def _timestamp() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat()


def _resolve_store_path(persistent_dir: str | Path, filename: str) -> Path:
    root = Path(persistent_dir)
    try:
        root.mkdir(parents=True, exist_ok=True)
        return root / filename
    except PermissionError:
        fallback_root = Path(tempfile.gettempdir()) / STORE_FALLBACK_DIRNAME
        fallback_root.mkdir(parents=True, exist_ok=True)
        return fallback_root / filename


def _load_store(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {}
    return payload if isinstance(payload, dict) else {}


def _save_store(path: Path, payload: dict[str, Any]) -> None:
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def _store_path(persistent_dir: str | Path, filename: str) -> Path:
    return _resolve_store_path(persistent_dir, filename)


def _owner_email(user_email: str | None) -> str:
    return (user_email or "private-space@maris.ai").strip().lower()


def _extract_preview(local_path: str, *, limit: int = 3) -> list[Any]:
    path = Path(local_path)
    if not path.is_file():
        return []
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return []
    if isinstance(payload, list):
        return payload[:limit]
    if isinstance(payload, dict):
        if isinstance(payload.get("preferences"), list):
            return payload["preferences"][:limit]
        if isinstance(payload.get("records"), list):
            return payload["records"][:limit]
        items = list(payload.items())[:limit]
        return [{"key": key, "value": value} for key, value in items]
    return []


def save_draft(
    persistent_dir: str | Path,
    *,
    user_email: str | None,
    name: str,
    payload: dict[str, Any],
    draft_id: str | None = None,
) -> dict[str, Any]:
    path = _store_path(persistent_dir, DRAFTS_FILENAME)
    owner = _owner_email(user_email)
    with STORE_LOCK:
        store = _load_store(path)
        drafts = store.setdefault("drafts", {})
        draft_id = (draft_id or uuid4().hex[:12]).strip()
        existing = drafts.get(draft_id, {})
        now = _timestamp()
        drafts[draft_id] = {
            "draft_id": draft_id,
            "name": name.strip(),
            "owner_email": owner,
            "payload": payload,
            "archived": False,
            "created_at": existing.get("created_at", now),
            "updated_at": now,
        }
        _save_store(path, store)
        return drafts[draft_id]


def list_drafts(
    persistent_dir: str | Path,
    *,
    user_email: str | None,
    include_archived: bool = False,
) -> list[dict[str, Any]]:
    path = _store_path(persistent_dir, DRAFTS_FILENAME)
    owner = _owner_email(user_email)
    with STORE_LOCK:
        drafts = list(_load_store(path).get("drafts", {}).values())
    items = [
        draft
        for draft in drafts
        if draft.get("owner_email") == owner and (include_archived or not draft.get("archived"))
    ]
    return sorted(items, key=lambda item: str(item.get("updated_at", "")), reverse=True)


def get_draft(persistent_dir: str | Path, draft_id: str) -> dict[str, Any] | None:
    path = _store_path(persistent_dir, DRAFTS_FILENAME)
    with STORE_LOCK:
        draft = _load_store(path).get("drafts", {}).get(draft_id)
    return draft if isinstance(draft, dict) else None


def archive_draft(
    persistent_dir: str | Path,
    *,
    draft_id: str,
    user_email: str | None,
) -> dict[str, Any] | None:
    path = _store_path(persistent_dir, DRAFTS_FILENAME)
    owner = _owner_email(user_email)
    with STORE_LOCK:
        store = _load_store(path)
        draft = store.get("drafts", {}).get(draft_id)
        if not isinstance(draft, dict) or draft.get("owner_email") != owner:
            return None
        draft["archived"] = True
        draft["updated_at"] = _timestamp()
        _save_store(path, store)
        return draft


def save_run(
    persistent_dir: str | Path,
    *,
    manifest: dict[str, Any],
    user_email: str | None,
    draft_id: str | None = None,
) -> dict[str, Any]:
    path = _store_path(persistent_dir, RUNS_FILENAME)
    owner = _owner_email(user_email)
    run_id = str(manifest["run_id"])
    now = _timestamp()
    record = {
        "run_id": run_id,
        "owner_email": owner,
        "draft_id": draft_id or "",
        "dataset_repo": manifest.get("dataset_repo", ""),
        "model_repo": manifest.get("model_repo", ""),
        "status": "staged",
        "ready_for_review": bool(manifest.get("ready_for_review")),
        "ready_for_training": bool(manifest.get("ready_for_training")),
        "quality_report": manifest.get("quality_report", {}),
        "input_summary": manifest.get("input_summary", {}),
        "artifact_count": len(manifest.get("artifacts", {})),
        "manifest": manifest,
        "created_at": now,
        "updated_at": now,
        "published_at": None,
        "training_started_at": None,
        "finished_at": None,
        "exit_code": None,
    }
    with STORE_LOCK:
        store = _load_store(path)
        runs = store.setdefault("runs", {})
        existing = runs.get(run_id)
        if isinstance(existing, dict):
            record["created_at"] = existing.get("created_at", now)
        runs[run_id] = record
        _save_store(path, store)
    return record


def get_run(persistent_dir: str | Path, run_id: str) -> dict[str, Any] | None:
    path = _store_path(persistent_dir, RUNS_FILENAME)
    with STORE_LOCK:
        record = _load_store(path).get("runs", {}).get(run_id)
    return record if isinstance(record, dict) else None


def list_runs(
    persistent_dir: str | Path,
    *,
    user_email: str | None,
    limit: int = 20,
) -> list[dict[str, Any]]:
    path = _store_path(persistent_dir, RUNS_FILENAME)
    owner = _owner_email(user_email)
    with STORE_LOCK:
        runs = list(_load_store(path).get("runs", {}).values())
    filtered = [item for item in runs if item.get("owner_email") == owner]
    return sorted(filtered, key=lambda item: str(item.get("updated_at", "")), reverse=True)[:limit]


def update_run(
    persistent_dir: str | Path,
    *,
    run_id: str,
    **fields: Any,
) -> dict[str, Any] | None:
    path = _store_path(persistent_dir, RUNS_FILENAME)
    with STORE_LOCK:
        store = _load_store(path)
        run = store.get("runs", {}).get(run_id)
        if not isinstance(run, dict):
            return None
        run.update(fields)
        run["updated_at"] = _timestamp()
        _save_store(path, store)
        return run


def index_artifacts(
    persistent_dir: str | Path,
    *,
    manifest: dict[str, Any],
    user_email: str | None,
    published: list[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
    path = _store_path(persistent_dir, ARTIFACTS_FILENAME)
    owner = _owner_email(user_email)
    run_id = str(manifest["run_id"])
    published_by_name = {
        str(item.get("artifact")): item
        for item in (published or [])
        if isinstance(item, dict) and item.get("artifact")
    }
    indexed: list[dict[str, Any]] = []
    with STORE_LOCK:
        store = _load_store(path)
        artifacts = store.setdefault("artifacts", {})
        for artifact_name, artifact in manifest.get("artifacts", {}).items():
            artifact_id = f"{run_id}:{artifact_name}"
            publish_meta = published_by_name.get(artifact_name, {})
            record = {
                "artifact_id": artifact_id,
                "artifact_name": artifact_name,
                "run_id": run_id,
                "owner_email": owner,
                "record_count": artifact.get("record_count", 0),
                "local_path": artifact.get("local_path", ""),
                "repo_path": publish_meta.get("path") or artifact.get("repo_path", ""),
                "published": bool(publish_meta),
                "preview": _extract_preview(str(artifact.get("local_path", ""))),
                "updated_at": _timestamp(),
            }
            existing = artifacts.get(artifact_id)
            if isinstance(existing, dict):
                record["created_at"] = existing.get("created_at", record["updated_at"])
                if existing.get("published") and not publish_meta:
                    record["published"] = True
                    record["repo_path"] = existing.get("repo_path", record["repo_path"])
            else:
                record["created_at"] = record["updated_at"]
            artifacts[artifact_id] = record
            indexed.append(record)
        _save_store(path, store)
    return indexed


def list_artifacts(
    persistent_dir: str | Path,
    *,
    user_email: str | None,
    run_id: str | None = None,
    limit: int = 50,
) -> list[dict[str, Any]]:
    path = _store_path(persistent_dir, ARTIFACTS_FILENAME)
    owner = _owner_email(user_email)
    with STORE_LOCK:
        items = list(_load_store(path).get("artifacts", {}).values())
    filtered = [
        item
        for item in items
        if item.get("owner_email") == owner and (not run_id or item.get("run_id") == run_id)
    ]
    return sorted(filtered, key=lambda item: str(item.get("updated_at", "")), reverse=True)[:limit]


def get_artifact(persistent_dir: str | Path, artifact_id: str) -> dict[str, Any] | None:
    path = _store_path(persistent_dir, ARTIFACTS_FILENAME)
    with STORE_LOCK:
        item = _load_store(path).get("artifacts", {}).get(artifact_id)
    return item if isinstance(item, dict) else None