import gradio as gr
from huggingface_hub import create_repo, upload_file, whoami, Repository
import os
import shutil

def duplicate(source_repo, dst_repo, token, repo_type, progress=gr.Progress()):
    # Get username from token
    username = whoami(token=token)["name"]
    
    progress(0, desc="Starting duplication process")
    logs = ["Starting duplication process..."]
    
    # Create the destination repo
    progress(0.1, desc="Creating destination repository")
    logs.append("Creating destination repository")
    if repo_type in ["space", "dataset"]:
        url = create_repo(
            repo_id=dst_repo,
            token=token,
            repo_type=repo_type,
            space_sdk="gradio" if repo_type == "space" else None,
            private=False
        )
    else:
        url = create_repo(
            repo_id=dst_repo,
            token=token,
            private=False
        )
    
    # Clone source repo
    progress(0.2, desc="Cloning source repository")
    logs.append("Cloning source repository")
    endpoint = "huggingface.co/"
    if repo_type in ["space", "dataset"]:
        endpoint += f"{repo_type}/"
    full_path = f"https://{username}:{token}@{endpoint}{source_repo}"
    local_dir = f"hub/{source_repo}"
    
    repo = Repository(
        local_dir=local_dir,
        clone_from=full_path,
        repo_type=repo_type if repo_type in ["space", "dataset"] else None,
        token=token
    )
    
    # Get list of files to upload
    files_to_upload = []
    for root, _, files in os.walk(local_dir):
        if not root.startswith(".") and ".git" not in root:
            for f in files:
                if not f.startswith("."):
                    files_to_upload.append((root, f))
    
    # Upload files with individual progress
    progress(0.3, desc="Preparing to upload files")
    logs.append(f"Found {len(files_to_upload)} files to upload")
    
    for i, (root, f) in enumerate(files_to_upload):
        file_progress = (i + 1) / len(files_to_upload)
        overall_progress = 0.3 + (0.6 * file_progress)  # 0.3 to 0.9 range for upload
        directory_path_in_repo = "/".join(root.split("/")[2:])
        path_in_repo = os.path.join(directory_path_in_repo, f)
        local_file_path = os.path.join(local_dir, path_in_repo)
        
        progress(
            (overall_progress, [
                (file_progress, f"Uploading {f}"),
            ]),
            desc=f"Uploading file {i+1}/{len(files_to_upload)}"
        )
        logs.append(f"Uploaded {f} to {path_in_repo}")
        
        upload_file(
            path_or_fileobj=local_file_path,
            path_in_repo=path_in_repo,
            repo_id=dst_repo,
            token=token,
            repo_type=repo_type if repo_type != "model" else None
        )
    
    # Clean up
    progress(0.95, desc="Cleaning up temporary files")
    logs.append("Cleaning up temporary files")
    shutil.rmtree(local_dir, ignore_errors=True)
    
    progress(1.0, desc="Completed")
    logs.append("Duplication completed successfully")
    
    return (
        f"Find your repo <a href='{url}' target='_blank' style='text-decoration:underline'>here</a>",
        "\n".join(logs)
    )

# Custom CSS for yellow progress bars
css = """
.progress-container:nth-child(4) .progress-bar {
    background-color: yellow !important;
}
"""

# Updated Gradio interface with progress panel
with gr.Blocks(title="Duplicate your repo!", css=css) as interface:
    gr.Markdown("""
    # Duplicate your repo!
    Duplicate a Hugging Face repository! You need a write token from https://hf.co/settings/tokens.
    This Space is an experimental demo.
    """)
    
    with gr.Row():
        # Left panel - Inputs
        with gr.Column(scale=2):
            source_input = gr.Textbox(
                label="Source repository",
                placeholder="e.g. osanseviero/src"
            )
            dest_input = gr.Textbox(
                label="Destination repository",
                placeholder="e.g. osanseviero/dst"
            )
            token_input = gr.Textbox(
                label="Write access token",
                type="password",
                placeholder="Your HF token"
            )
            repo_type_input = gr.Dropdown(
                choices=["model", "dataset", "space"],
                label="Repository type",
                value="model"
            )
            submit_btn = gr.Button("Duplicate")
            output = gr.HTML(label="Result")
        
        # Right panel - Progress
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("## Progress")
                gr.Markdown("### Overall Progress")
                global_progress = gr.Progress()
                gr.Markdown("### File Progress")
                with gr.Column(variant="panel", elem_classes="progress-container"):
                    file_progress = gr.Progress()
                log_output = gr.Textbox(label="Progress Log", lines=10)
    
    gr.Markdown("""
    Find your write token at 
    [token settings](https://huggingface.co/settings/tokens)
    """)
    
    submit_btn.click(
        fn=duplicate,
        inputs=[source_input, dest_input, token_input, repo_type_input],
        outputs=[output, log_output]
    )

interface.launch()