#!/usr/bin/env python3
"""Restore data from HF Dataset snapshot before server starts.

Handles three cases:
  1. PAPERCLIP_DATA_REPO_ID or HF_TOKEN not set -> skip (exit 0)
  2. Dataset repo does not exist (first deploy) -> skip gracefully (exit 0)
  3. Restore succeeds -> files copied to PAPERCLIP_HOME (exit 0)
  4. Unexpected error -> FATAL, exit 1 (refuse to start)
"""
import os
import shutil
import sys
import tempfile

from huggingface_hub import snapshot_download
from huggingface_hub.errors import RepositoryNotFoundError


def _get_env(name, default=""):
    """Read env var at call time (not import time) so tests can monkeypatch."""
    return os.environ.get(name, default)


def _get_instance_name():
    """Derive instance name from SPACE_ID or fallback env var.

    SPACE_ID on HF Spaces is like 'user/space-name'.
    We sanitize '/' to '-' for use as a subdirectory name.
    """
    name = (
        _get_env("SPACE_ID")
        or _get_env("PAPERCLIP_INSTANCE_NAME")
        or _get_env("PAPERCLIP_INSTANCE_ID", "default")
    )
    return name.replace("/", "-")


def _sync_restore(src_dir: str, dst_dir: str) -> None:
    """Restore snapshot using atomic file replacement.

    For files: write to a temp path next to the target, then os.replace()
    (atomic on same filesystem). If the process crashes mid-restore,
    old files remain intact and temp files are harmless.

    For directories: recurse depth-first so leaf files are replaced atomically.
    Local-only entries (not in snapshot) are left untouched.
    """
    for entry in os.listdir(src_dir):
        src_path = os.path.join(src_dir, entry)
        dst_path = os.path.join(dst_dir, entry)
        if os.path.isdir(src_path):
            os.makedirs(dst_path, exist_ok=True)
            _sync_restore(src_path, dst_path)
        else:
            os.makedirs(dst_dir, exist_ok=True)
            tmp_path = dst_path + ".restoring"
            shutil.copy2(src_path, tmp_path)
            os.replace(tmp_path, dst_path)


def main():
    paperclip_home = _get_env("PAPERCLIP_HOME", "/paperclip")
    data_repo_id = _get_env("PAPERCLIP_DATA_REPO_ID")
    hf_token = _get_env("HF_TOKEN") or _get_env("HUGGING_FACE_HUB_TOKEN")

    if not data_repo_id:
        print("restore_snapshot: PAPERCLIP_DATA_REPO_ID not set, skipping restore")
        return

    if not hf_token:
        print("restore_snapshot: HF_TOKEN not set, skipping restore")
        return

    instance_name = _get_instance_name()
    print(f"restore_snapshot: restoring {instance_name} from {data_repo_id}")

    try:
        with tempfile.TemporaryDirectory() as tmp_dir:
            snapshot_download(
                repo_id=data_repo_id,
                repo_type="dataset",
                local_dir=tmp_dir,
                token=hf_token,
                allow_patterns=[f"{instance_name}/**"],
            )
            # snapshot_download preserves subdirectory structure, so
            # files land at tmp_dir/{instance_name}/...
            instance_dir = os.path.join(tmp_dir, instance_name)
            if os.path.isdir(instance_dir):
                _sync_restore(instance_dir, paperclip_home)
                print(f"restore_snapshot: restored data to {paperclip_home}")
            else:
                print(
                    f"restore_snapshot: no data for instance {instance_name}, skipping"
                )

    except RepositoryNotFoundError:
        print(
            f"restore_snapshot: Dataset {data_repo_id} not found "
            "(first deploy?), skipping restore"
        )
    except Exception as e:
        print(f"FATAL: restore_snapshot failed: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()