import gradio as gr
from huggingface_hub import login, snapshot_download, HfApi
import os
import shutil

def download_and_upload(token, source_repo_id, repo_type, target_repo_id, path_in_repo=""):
    try:
        # Use HF_TOKEN secret from environment if input is empty
        if token.strip() == "":
            token = os.environ.get("HF_TOKEN")
        if not token:
            return "❌ Token missing. Provide one or set HF_TOKEN in secrets."

        login(token=token)
        api = HfApi()

        # Try to create target repo (will skip if already exists)
        api.create_repo(repo_id=target_repo_id, repo_type=repo_type, exist_ok=True)

        # Download
        if repo_type == "model":
            path = snapshot_download(repo_id=source_repo_id)
        else:
            path = snapshot_download(repo_id=source_repo_id, repo_type=repo_type)

        # Upload
        api.upload_folder(
            folder_path=path,
            repo_id=target_repo_id,
            repo_type=repo_type,
            path_in_repo=path_in_repo,
        )

        return f"✅ Downloaded from {source_repo_id} and uploaded to {target_repo_id} (created if not exist)!"

    except Exception as e:
        return f"❌ Failed: {e}"


with gr.Blocks() as demo:
    gr.Markdown("## 🤗 Auto Download & Upload HF Repo (With Auto-Cleanup)")

    token = gr.Textbox(label="Hugging Face Token (leave blank to use secret)", type="password", placeholder="hf_...")
    source_repo_id = gr.Textbox(label="Source Repo ID (e.g. user/source-model)")
    repo_type = gr.Radio(choices=["model", "dataset", "space"], label="Repo Type", value="model")
    target_repo_id = gr.Textbox(label="Target Repo ID (e.g. user/target-model)")
    path_in_repo = gr.Textbox(label="Path in Target Repo (optional)", value="", placeholder="e.g. subfolder or leave empty")

    output = gr.Textbox(label="Result")

    run_button = gr.Button("Start Download + Upload")
    run_button.click(
        fn=download_and_upload,
        inputs=[token, source_repo_id, repo_type, target_repo_id, path_in_repo],
        outputs=output,
    )

demo.launch()