#!/usr/bin/env python3
"""Restore Open WebUI data into the data dir from the backups dataset.

Runs on every boot. If webui.db is already present the volume is intact and
nothing happens. If it is missing (volume wiped / factory reset), the latest
full archive is extracted into the data dir, falling back to the legacy
single-file DB backup written by older versions of backup.sh.

Never fails the boot: a restore error is logged and ignored, so a first boot
with no backup yet still starts the app.

Set OWUI_RESTORE_ARCHIVE to restore a pinned snapshot instead of the rolling
latest, e.g. backups/daily/open-webui-data-2026-10-28.tar.gz (list what exists
with HfApi().list_repo_files('tkgray/open-webui-backups', repo_type='dataset')).
"""
import os
import shutil
import sys
import tarfile

DATA_DIR = os.environ.get("OWUI_DATA_DIR", "/data")
REPO_ID = os.environ.get("OWUI_BACKUP_REPO", "tkgray/open-webui-backups")
ARCHIVE_IN_REPO = os.environ.get("OWUI_RESTORE_ARCHIVE", "backups/open-webui-data.tar.gz")
LEGACY_DB_IN_REPO = "backups/webui.db"

DB_PATH = os.path.join(DATA_DIR, "webui.db")


def log(msg):
    print(f"[restore] {msg}", flush=True)


def fetch(token, filename):
    from huggingface_hub import hf_hub_download

    return hf_hub_download(
        repo_id=REPO_ID,
        filename=filename,
        repo_type="dataset",
        token=token,
    )


def safe_extract(tf, dest):
    """Extract an archive into dest, refusing absolute paths and .. traversal."""
    dest = os.path.abspath(dest)
    for member in tf.getmembers():
        target = os.path.abspath(os.path.join(dest, member.name))
        if target != dest and not target.startswith(dest + os.sep):
            raise ValueError(f"unsafe path in archive: {member.name}")
    try:
        tf.extractall(dest, filter="data")  # py3.12+
    except TypeError:
        tf.extractall(dest)


def main():
    if os.path.isfile(DB_PATH):
        log(f"{DB_PATH} already exists — nothing to restore.")
        return 0

    token = os.environ.get("HF_TOKEN")
    if not token:
        log("HF_TOKEN not set — skipping restore.")
        return 0

    os.makedirs(DATA_DIR, exist_ok=True)
    log(f"{DB_PATH} missing — attempting restore from {REPO_ID} ...")

    # Preferred: the full archive (DB + vector_db + uploads)
    try:
        path = fetch(token, ARCHIVE_IN_REPO)
        with tarfile.open(path, "r:gz") as tf:
            safe_extract(tf, DATA_DIR)
        if os.path.isfile(DB_PATH):
            log("Full archive restore OK")
            return 0
        log("archive extracted but webui.db is missing — falling back")
    except Exception as exc:  # noqa: BLE001
        log(f"archive restore failed: {exc}")

    # Fallback: legacy single-file DB backup
    try:
        path = fetch(token, LEGACY_DB_IN_REPO)
        shutil.copy2(path, DB_PATH)
        log("Legacy DB restore OK")
    except Exception as exc:  # noqa: BLE001
        log(f"no backup available yet: {exc}")

    return 0


if __name__ == "__main__":
    sys.exit(main())