import os
import tarfile
import time
from huggingface_hub import hf_hub_download, upload_file

DATASET_REPO = "cys003/chromium-config"   # ← 确认是这个
FILE_NAME = "config.tar.gz"
CONFIG_DIR = "/config"
TOKEN = os.environ.get("HF_TOKEN")

def restore():
    print("=== 开始执行 Chromium 配置恢复 ===")
    try:
        local_file = hf_hub_download(
            repo_id=DATASET_REPO,
            filename=FILE_NAME,
            repo_type="dataset",
            token=TOKEN
        )
        os.makedirs(CONFIG_DIR, exist_ok=True)
        with tarfile.open(local_file, "r:gz") as tar:
            tar.extractall(CONFIG_DIR, filter='data')
        
        # 清理锁文件
        for lock_file in ["SingletonLock", "SingletonSocket", "SingletonCookie", ".lock"]:
            lock_path = os.path.join(CONFIG_DIR, lock_file)
            if os.path.exists(lock_path):
                os.remove(lock_path)
                print(f"已清理锁文件: {lock_file}")
        
        print("=== 配置恢复成功！ ===")
    except Exception as e:
        print("首次启动或无备份，跳过恢复:", str(e))

def backup():
    print("=== 开始备份 Chromium 配置 ===")
    try:
        tar_path = "/tmp/config.tar.gz"
        with tarfile.open(tar_path, "w:gz") as tar:
            tar.add(CONFIG_DIR, arcname=".")
        upload_file(
            path_or_fileobj=tar_path,
            path_in_repo=FILE_NAME,
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=TOKEN,
            commit_message=f"Auto backup {time.strftime('%Y-%m-%d %H:%M:%S')}"
        )
        print("=== 备份上传成功！ ===")
        os.remove(tar_path)
    except Exception as e:
        print("备份失败:", str(e))

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