import gradio as gr
import os
import shutil
from pathlib import Path
from huggingface_hub import HfApi, hf_hub_download

# --- 持久化类 (直接嵌入以便于 Demo 运行) ---
class PersistenceManager:
    def __init__(self, dataset_id, token=None):
        self.api = HfApi(token=token or os.environ.get("HF_TOKEN"))
        self.dataset_id = dataset_id

    def restore(self, remote_path, local_path):
        try:
            downloaded_path = hf_hub_download(
                repo_id=self.dataset_id,
                repo_type="dataset",
                filename=remote_path,
                token=os.environ.get("HF_TOKEN")
            )
            shutil.copy(downloaded_path, local_path)
            return True
        except Exception:
            return False

    def save(self, local_path, remote_path):
        try:
            self.api.upload_file(
                path_or_fileobj=local_path,
                path_in_repo=remote_path,
                repo_id=self.dataset_id,
                repo_type="dataset",
                token=os.environ.get("HF_TOKEN")
            )
            return True
        except Exception:
            return False

# --- App 逻辑 ---
DATASET_ID = os.environ.get("DATASET_REPO_ID", "mingyang22/persistent-storage")
DB_FILE = "counter.txt"
REMOTE_PATH = "backups/counter.txt"

pm = PersistenceManager(DATASET_ID)

# 初始化：尝试恢复
pm.restore(REMOTE_PATH, DB_FILE)

def get_count():
    if os.path.exists(DB_FILE):
        with open(DB_FILE, "r") as f:
            return f.read().strip()
    return "0"

def increment():
    current = int(get_count())
    new_count = current + 1
    with open(DB_FILE, "w") as f:
        f.write(str(new_count))
    
    # 异步备份 (为了快，这里同步写，但在生产中可以放到后台)
    success = pm.save(DB_FILE, REMOTE_PATH)
    status = "✅ 备份成功" if success else "❌ 备份失败 (请确认 Secret 配置)"
    return str(new_count), status

with gr.Blocks() as demo:
    gr.Markdown("# 🚀 HuggingFace 持久化 Demo")
    gr.Markdown("这个 Space 在重启后会通过 Dataset 恢复之前的计数状态。")
    
    with gr.Row():
        counter_display = gr.Label(value=get_count(), label="当前计数值")
        status_display = gr.Textbox(label="云端同步状态")
    
    btn = gr.Button("累加并存档 (Increment & Save)")
    btn.click(fn=increment, outputs=[counter_display, status_display])
    
    gr.Markdown("### 如何验证？")
    gr.Markdown("1. 点击按钮累加数字。\n2. 在管理脚本中执行 `restart` 命令重启这个 Space。\n3. 刷新页面，你会发现计数值没有归零，而是从上次保存的地方继续！")

demo.launch()