import os
import subprocess
import tempfile
import shutil

BACKUP_REPO = os.environ.get("BACKUP_REPO")
HF_TOKEN = os.environ.get("HF_TOKEN")

if not BACKUP_REPO or not HF_TOKEN:
    print("[Restore] Skipping: BACKUP_REPO or HF_TOKEN not set")
    exit(0)

env = os.environ.copy()
env["HF_HOME"] = "/tmp/hf_cache"
env["XDG_CACHE_HOME"] = "/tmp/xdg_cache"
env["TMPDIR"] = "/tmp"
env["HF_TOKEN"] = HF_TOKEN

os.makedirs(env["HF_HOME"], exist_ok=True)
os.makedirs(env["XDG_CACHE_HOME"], exist_ok=True)
os.makedirs(env["TMPDIR"], exist_ok=True)

# Step 1: List repo files
print("[Restore] Listing backups in repo...")
result = subprocess.run(
    ["hf", "ls-files", BACKUP_REPO, "--repo-type", "dataset"],
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True,
    env=env,
)
files = result.stdout.strip().splitlines()
backups = [f for f in files if f.endswith(".tar.gz")]

if not backups:
    print("[Restore] No backups found")
    exit(0)

# Step 2: Pick latest (they’re timestamped)
latest = sorted(backups)[-1]
print(f"[Restore] Found latest backup: {latest}")

# Step 3: Download
tmpfile = os.path.join(tempfile.gettempdir(), "restore.tar.gz")
subprocess.run(
    ["hf", "download", BACKUP_REPO, latest, "--repo-type", "dataset", "--local-dir", "/tmp/restore", "--force"],
    check=True,
    env=env,
)

downloaded = os.path.join("/tmp/restore", latest)

# Step 4: Extract into /home/vscode
print("[Restore] Extracting backup...")
if os.path.exists("/home/vscode"):
    shutil.rmtree("/home/vscode")
os.makedirs("/home/vscode", exist_ok=True)

subprocess.run(
    ["tar", "-xzf", downloaded, "-C", "/"],
    check=True
)

print("[Restore] Completed")