#!/usr/bin/env python3
"""Back up Open WebUI user data to a Hugging Face dataset repo.

Covers everything a user would miss after a volume wipe:
  * webui.db        - users, logins, settings, chats (consistent SQLite snapshot)
  * vector_db/      - RAG embeddings for uploaded documents
  * uploads/        - user-uploaded files

Deliberately EXCLUDES /data/cache (the embedding-model cache, ~1 GB, regenerates
itself on boot).

Two copies go up per run: the rolling latest at backups/open-webui-data.tar.gz
(what restore.py reads on boot) and a dated snapshot
backups/daily/open-webui-data-<YYYY-MM-DD>.tar.gz. The dated snapshots are what
make a bad wipe survivable - the rolling file is overwritten every hour, so a
corrupted dataset gets baked into it within one interval. The newest
OWUI_BACKUP_DAILY_KEEP (default 7) dated snapshots are kept, older ones deleted.

Requires the HF_TOKEN env var (Space secret). Never fails the app.
"""
import datetime
import gzip
import os
import shutil
import sqlite3
import sys
import tarfile
import tempfile

DATA_DIR = os.environ.get("OWUI_DATA_DIR", "/data")
REPO_ID = os.environ.get("OWUI_BACKUP_REPO", "tkgray/open-webui-backups")
ARCHIVE_IN_REPO = "backups/open-webui-data.tar.gz"
DAILY_DIR = "backups/daily"
DAILY_PREFIX = "open-webui-data-"
INCLUDE_DIRS = ("vector_db", "uploads")


def env_int(name, default):
    raw = os.environ.get(name)
    if raw is None:
        return default
    try:
        return int(raw)
    except (TypeError, ValueError):
        return default


DAILY_KEEP = env_int("OWUI_BACKUP_DAILY_KEEP", 7)

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


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


def snapshot_sqlite(src, dst):
    """Write a consistent copy of a live SQLite DB using the backup API.

    Copying the file directly (or tarring it) can capture a torn WAL, so use
    sqlite3's own online-backup, which is safe while the app is writing.
    """
    source = sqlite3.connect(f"file:{src}?mode=ro", uri=True)
    try:
        dest = sqlite3.connect(dst)
        try:
            with dest:
                source.backup(dest)
        finally:
            dest.close()
    finally:
        source.close()


def build_archive(payload_dir, archive_path):
    """Deterministically tar.gz a staged payload.

    Every member's mtime/uid/gid and the gzip header mtime are zeroed and the
    member order is fixed, so an unchanged payload yields byte-identical output.
    That lets the dataset host recognise a no-op upload and skip the commit,
    instead of creating a fresh commit every hour.
    """
    entries = []
    for root, _dirs, files in os.walk(payload_dir):
        for name in files:
            full = os.path.join(root, name)
            entries.append((os.path.relpath(full, payload_dir), full))
    entries.sort()

    with open(archive_path, "wb") as raw:
        # filename="" suppresses the gzip FNAME header, which would otherwise
        # carry the temp archive path and make every run's bytes differ.
        with gzip.GzipFile(filename="", fileobj=raw, mode="wb", mtime=0) as gz:
            with tarfile.open(fileobj=gz, mode="w") as tf:
                for rel, full in entries:
                    info = tf.gettarinfo(full, arcname=rel)
                    info.mtime = 0
                    info.uid = 0
                    info.gid = 0
                    info.uname = ""
                    info.gname = ""
                    with open(full, "rb") as fh:
                        tf.addfile(info, fh)
    return archive_path


def prune_daily(api, token):
    """Delete dated snapshots beyond the newest DAILY_KEEP.

    Filenames are ISO-dated, so a plain sort is chronological. Only ever
    deletes under DAILY_DIR, so the rolling archive and the legacy single-file
    backup are untouchable by this path.
    """
    try:
        files = api.list_repo_files(REPO_ID, repo_type="dataset", token=token)
    except Exception as exc:  # noqa: BLE001 - a prune failure must not fail the run
        log(f"prune skipped, listing failed: {exc}")
        return

    dated = sorted(
        f for f in files
        if f.startswith(f"{DAILY_DIR}/{DAILY_PREFIX}") and f.endswith(".tar.gz")
    )
    stale = dated[:-DAILY_KEEP] if DAILY_KEEP > 0 else dated
    if not stale:
        log(f"prune: {len(dated)} dated snapshot(s) on disk, limit {DAILY_KEEP} - nothing to delete")
        return

    from huggingface_hub import CommitOperationDelete, create_commit

    create_commit(
        repo_id=REPO_ID,
        repo_type="dataset",
        operations=[CommitOperationDelete(path_in_repo=path) for path in stale],
        commit_message=f"prune {len(stale)} snapshot(s) beyond {DAILY_KEEP}-day retention",
        token=token,
    )
    log(f"prune: deleted {', '.join(stale)}")


def main():
    if not os.path.isfile(DB_PATH):
        log(f"{DB_PATH} does not exist yet — skipping.")
        return 0

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

    stage = tempfile.mkdtemp()
    payload = os.path.join(stage, "data")
    os.makedirs(payload, exist_ok=True)
    try:
        snapshot_sqlite(DB_PATH, os.path.join(payload, "webui.db"))
        for name in INCLUDE_DIRS:
            src = os.path.join(DATA_DIR, name)
            if os.path.isdir(src):
                shutil.copytree(src, os.path.join(payload, name), symlinks=True)

        archive = os.path.join(stage, "open-webui-data.tar.gz")
        build_archive(payload, archive)
        log(f"staged {os.path.getsize(archive)} bytes; uploading to {REPO_ID} ...")

        from huggingface_hub import HfApi

        api = HfApi(token=token)
        dated_in_repo = f"{DAILY_DIR}/{DAILY_PREFIX}{datetime.date.today().isoformat()}.tar.gz"
        for path_in_repo in (ARCHIVE_IN_REPO, dated_in_repo):
            api.upload_file(
                path_or_fileobj=archive,
                path_in_repo=path_in_repo,
                repo_id=REPO_ID,
                repo_type="dataset",
                token=token,
            )
            log(f"Upload OK -> {path_in_repo}")

        prune_daily(api, token)
        return 0
    except Exception as exc:  # noqa: BLE001 - never fail the app over a backup
        log(f"Upload FAILED: {exc}")
        return 1
    finally:
        shutil.rmtree(stage, ignore_errors=True)


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