import os
import time
import schedule
from huggingface_hub import HfApi, hf_hub_download
import subprocess
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = os.getenv("DATASET_REPO", "augment17/agentscope-redis-backup")
DUMP_FILE = "/data/dump.rdb"
LOCAL_DUMP = "dump.rdb"

api = HfApi(token=HF_TOKEN)

def ensure_repo():
    try:
        api.repo_info(repo_id=DATASET_REPO, repo_type="dataset")
    except Exception:
        logger.info(f"Creating dataset repo {DATASET_REPO}...")
        api.create_repo(repo_id=DATASET_REPO, repo_type="dataset", private=True)

def restore_backup():
    ensure_repo()
    try:
        logger.info("Checking for existing Redis backup...")
        downloaded_file = hf_hub_download(
            repo_id=DATASET_REPO, 
            repo_type="dataset", 
            filename=LOCAL_DUMP,
            token=HF_TOKEN,
            local_dir="/data"
        )
        logger.info(f"Restored backup from {downloaded_file}")
    except Exception as e:
        logger.info(f"No existing backup found or failed to restore: {e}")

def perform_backup():
    logger.info("Performing Redis backup...")
    try:
        # Trigger Redis bgsave
        subprocess.run(["redis-cli", "save"], check=True)
        if os.path.exists(DUMP_FILE):
            api.upload_file(
                path_or_fileobj=DUMP_FILE,
                path_in_repo=LOCAL_DUMP,
                repo_id=DATASET_REPO,
                repo_type="dataset",
                token=HF_TOKEN
            )
            logger.info("Backup uploaded to Hugging Face Datasets successfully.")
        else:
            logger.error("dump.rdb not found after save!")
    except Exception as e:
        logger.error(f"Backup failed: {e}")

if __name__ == "__main__":
    mode = os.getenv("MODE", "restore")
    if mode == "restore":
        restore_backup()
    elif mode == "cron":
        logger.info("Starting backup cron schedule (every 30 minutes)...")
        schedule.every(30).minutes.do(perform_backup)
        while True:
            schedule.run_pending()
            time.sleep(60)