import os
import sys
import tarfile
import time
from huggingface_hub import HfApi, hf_hub_download

api = HfApi()
repo_id = os.getenv("HF_DATASET")
token = os.getenv("HF_TOKEN")
FILENAME = "latest_backup.tar.gz"

# ===== 备份配置 =====
# 跳过最近多少秒内被修改的 session 文件(避免 tar 读到 agent 正在写的文件)
ACTIVE_SESSION_THRESHOLD_SECONDS = 1800  # 30 分钟

# 备份根目录(用于把 arcname 还原成绝对路径,做 mtime 检查)
BACKUP_ROOT = "/root/.openclaw"

# 要备份的子路径(相对 BACKUP_ROOT)
PATHS_TO_BACKUP = [
    "sessions",
    "agents/main/sessions",
    "openclaw.json",
    "vscode_cli_data",
    "cron",
    "credentials",
    "skills",
    "plugins",
    "memory",
    "plugin-skills",
]



def restore():
    try:
        if not repo_id or not token:
            print("Skip Restore: HF_DATASET or HF_TOKEN not set")
            return
        
        # 直接下载最新文件
        print(f"Downloading {FILENAME} from {repo_id}...")
        path = hf_hub_download(repo_id=repo_id, filename=FILENAME, repo_type="dataset", token=token)
        
        with tarfile.open(path, "r:gz") as tar:
            tar.extractall(path="/root/.openclaw/")
        print(f"Success: Restored from {FILENAME}")
        return True
    except Exception as e:
        # 如果是第一次运行，仓库里没文件，报错是正常的
        print(f"Restore Note: No existing backup found or error: {e}")

def backup():
    try:
        if not repo_id or not token:
            print("Skip Backup: HF_DATASET or HF_TOKEN not set")
            return

        skipped_active = []  # 收集被跳过的 active session,最后打一行日志

        def exclude_filter(tarinfo):
            parts = tarinfo.name.split("/")

            # 规则 1:排除 skills/data(skill 输出/缓存目录,不需要备份)
            if len(parts) >= 2 and parts[0] == "skills" and parts[1] == "data":
                return None

            # 规则 2:跳过最近 30 分钟内被修改的 .jsonl session 文件
            # 这些文件 agent 可能正在写入,tar 一边读它一边被改,会触发
            # EmbeddedAttemptSessionTakeoverError(session file changed while lock released)
            if tarinfo.name.endswith(".jsonl"):
                full_path = os.path.join(BACKUP_ROOT, tarinfo.name)
                try:
                    age = time.time() - os.path.getmtime(full_path)
                    if age < ACTIVE_SESSION_THRESHOLD_SECONDS:
                        skipped_active.append(tarinfo.name)
                        return None
                except OSError:
                    # 文件读不到 mtime 就跳过,安全起见
                    return None

            return tarinfo

        with tarfile.open(FILENAME, "w:gz") as tar:
            for sub in PATHS_TO_BACKUP:
                full = os.path.join(BACKUP_ROOT, sub)
                if os.path.exists(full):
                    tar.add(full, arcname=sub, filter=exclude_filter)
                else:
                    print(f"Skip (not exist): {full}")

        if skipped_active:
            print(f"Skipped {len(skipped_active)} active session file(s) modified within "
                  f"{ACTIVE_SESSION_THRESHOLD_SECONDS}s")

        api.upload_file(
            path_or_fileobj=FILENAME,
            path_in_repo=FILENAME,
            repo_id=repo_id,
            repo_type="dataset",
            token=token
        )

        # 打印备份大小,方便观察
        size_mb = os.path.getsize(FILENAME) / 1024 / 1024
        print(f"Backup {FILENAME} Success ({size_mb:.2f} MB, Overwritten).")

    except Exception as e:
        print(f"Backup Error: {e}")
        sys.exit(1)

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