import os
import sys
import re
import tarfile
from datetime import datetime, timezone

from huggingface_hub import HfApi, hf_hub_download

api = HfApi()
repo_id = os.getenv("HF_DATASET")
token = os.getenv("HF_TOKEN")

OPENCLAW_DIR = "/root/.openclaw"
REPO_TYPE = "dataset"

# 备份文件命名：backup-YYYYmmdd-HHMMSS.tar.gz（带时间戳，便于轮转）
BACKUP_PREFIX = "backup-"
BACKUP_SUFFIX = ".tar.gz"
BACKUP_RE = re.compile(r"^backup-\d{8}-\d{6}\.tar\.gz$")
LEGACY_FILENAME = "latest_backup.tar.gz"  # 兼容历史上的单文件备份

KEEP = 5  # 仅保留最近 5 份备份，多余的删除

# 备份时跳过的临时文件/目录
EXCLUDE_SUFFIXES = (".log", ".tmp", ".bak", ".pid")
EXCLUDE_NAMES = {"__pycache__", "node_modules", ".git"}


def _check_env():
    if not repo_id or not token:
        print("Skip: HF_DATASET or HF_TOKEN not set")
        return False
    return True


def should_exclude(tarinfo):
    """过滤掉不需要备份的临时文件"""
    name = os.path.basename(tarinfo.name)
    if name in EXCLUDE_NAMES:
        return None
    if any(name.endswith(s) for s in EXCLUDE_SUFFIXES):
        return None
    return tarinfo


def _list_backups():
    """返回 (备份文件名列表[新->旧], 仓库全部文件列表)"""
    files = api.list_repo_files(repo_id=repo_id, repo_type=REPO_TYPE, token=token)
    backups = sorted((f for f in files if BACKUP_RE.match(f)), reverse=True)
    return backups, files


def restore():
    if not _check_env():
        return
    try:
        backups, files = _list_backups()
        # 优先恢复最新的时间戳备份；否则回退到旧的单文件备份
        target = backups[0] if backups else (
            LEGACY_FILENAME if LEGACY_FILENAME in files else None)
        if not target:
            print("Restore Note: no backup found in repo")
            return

        print(f"Downloading {target} from {repo_id}...")
        path = hf_hub_download(
            repo_id=repo_id, filename=target,
            repo_type=REPO_TYPE, token=token)

        with tarfile.open(path, "r:gz") as tar:
            tar.extractall(path=OPENCLAW_DIR)

        print(f"Success: restored from {target}")
        return True
    except Exception as e:
        print(f"Restore Note: no existing backup found or error: {e}")


def _prune_and_squash():
    """删除超过 KEEP 份的旧备份 + 清理遗留单文件，并压扁 git 历史真正回收存储。"""
    try:
        backups, files = _list_backups()

        # 1) 需要删除的：超过 KEEP 份的旧备份 + 遗留的 latest_backup.tar.gz
        to_delete = list(backups[KEEP:])
        if LEGACY_FILENAME in files and len(backups) >= 1:
            to_delete.append(LEGACY_FILENAME)

        for name in to_delete:
            try:
                api.delete_file(
                    path_in_repo=name, repo_id=repo_id,
                    repo_type=REPO_TYPE, token=token,
                    commit_message=f"prune old backup {name}")
                print(f"Pruned: {name}")
            except Exception as e:
                print(f"Prune warn ({name}): {e}")

        # 2) 关键一步：压扁全部历史，否则被删/被覆盖的旧版本 blob 仍占配额
        try:
            api.super_squash_history(
                repo_id=repo_id, repo_type=REPO_TYPE, token=token,
                commit_message="squash backup history to reclaim storage")
            print("Squashed repo history (storage reclaimed)")
        except Exception as e:
            print(f"Squash warn: {e}")
    except Exception as e:
        print(f"Prune/squash warn: {e}")


def backup():
    if not _check_env():
        return
    if not os.path.exists(OPENCLAW_DIR):
        print(f"Skip Backup: {OPENCLAW_DIR} does not exist")
        return

    ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    fname = f"{BACKUP_PREFIX}{ts}{BACKUP_SUFFIX}"
    uploaded = False
    try:
        with tarfile.open(fname, "w:gz") as tar:
            tar.add(OPENCLAW_DIR, arcname=".", filter=should_exclude)

        size_mb = os.path.getsize(fname) / (1024 * 1024)
        print(f"Backup archive size: {size_mb:.1f} MB ({fname})")

        api.upload_file(
            path_or_fileobj=fname, path_in_repo=fname,
            repo_id=repo_id, repo_type=REPO_TYPE, token=token)
        uploaded = True
        print(f"Backup {fname} uploaded.")
    except Exception as e:
        # 上传失败（如配额超限）必须显式告警，绝不能静默吞掉
        print(f"!!! BACKUP FAILED (check HF storage quota / token!): {e}")
    finally:
        try:
            os.remove(fname)
        except OSError:
            pass

    # 仅在成功上传后再做轮转 + 压扁历史
    if uploaded:
        _prune_and_squash()


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "backup":
        backup()
    else:
        restore()