import gradio as gr
import requests
from pathlib import Path
import re
import os
from huggingface_hub import whoami, HfApi, hf_hub_download
from huggingface_hub.utils import build_hf_headers, hf_raise_for_status
from gradio_huggingfacehub_search import HuggingfaceHubSearch

ENDPOINT = "https://huggingface.co"
# ENDPOINT = "http://localhost:5564"

REPO_TYPES = ["model", "dataset", "space"]
HF_REPO = os.environ.get("HF_REPO") if os.environ.get("HF_REPO") else "" # set your default repo


def duplicate(source_repo, dst_repo, repo_type, private, overwrite, auto_dir, oauth_token: gr.OAuthToken | None, progress=gr.Progress(track_tqdm=True)):
    print(oauth_token.token)
    hf_token = oauth_token.token
    api = HfApi(token=hf_token)
    try:
        if not repo_type in REPO_TYPES:
            raise ValueError("need to select valid repo type")
        _ = whoami(oauth_token.token)
        # ^ this will throw if token is invalid

        if re.fullmatch(r'^[\w_\-\.]+/[\w_\-\.]+$', dst_repo): subfolder = ""
        else:
            dst_repo, subfolder = re.findall(r'^([\w_\-\.]+/[\w_\-\.]+)/?(.+)?$', dst_repo)[0]
            subfolder = subfolder.removesuffix("/")
        if auto_dir: subfolder = source_repo

        if overwrite and api.repo_exists(repo_id=dst_repo, repo_type=repo_type, token=hf_token) or subfolder:
            api.create_repo(repo_id=dst_repo, repo_type=repo_type, private=private, exist_ok=True, token=hf_token)
            for path in api.list_repo_files(repo_id=source_repo, repo_type=repo_type, token=hf_token):
                file = hf_hub_download(repo_id=source_repo, filename=path, repo_type=repo_type, token=hf_token)
                if not Path(file).exists(): continue
                if Path(file).is_dir(): # unused for now
                    api.upload_folder(repo_id=dst_repo, folder_path=file, path_in_repo=f"{subfolder}/{path}" if subfolder else path, repo_type=repo_type, token=hf_token)
                elif Path(file).is_file():
                    api.upload_file(repo_id=dst_repo, path_or_fileobj=file, path_in_repo=f"{subfolder}/{path}" if subfolder else path, repo_type=repo_type, token=hf_token)
                if Path(file).exists(): Path(file).unlink()
            if repo_type == "dataset": repo_url = f"https://huggingface.co/datasets/{dst_repo}"
            elif repo_type == "space": repo_url = f"https://huggingface.co/spaces/{dst_repo}"
            else: repo_url = f"https://huggingface.co/{dst_repo}"
        else:
            r = requests.post(
                f"{ENDPOINT}/api/{repo_type}s/{source_repo}/duplicate",
                headers=build_hf_headers(token=oauth_token.token),
                json={"repository": dst_repo, "private": private},
            )
            hf_raise_for_status(r)

            repo_url = r.json().get("url")

        return (
            f'Find your repo <a href=\'{repo_url}\' target="_blank" style="text-decoration:underline">here</a>',
            "sp.jpg",
        )

    except Exception as e:
        raise gr.Error(f"""Oops, you forgot to login. Please use the loggin button on the top left to migrate your repo {e}""")

interface = gr.Interface(
    fn=duplicate,
    inputs=[
        HuggingfaceHubSearch(
            placeholder="Source repository (e.g. osanseviero/src)",
            search_type=["model", "dataset", "space"],
            sumbit_on_select=False,
        ),
        gr.Textbox(placeholder="Destination repository (e.g. osanseviero/dst)", value=HF_REPO),
        gr.Dropdown(choices=REPO_TYPES, value="model"),
        gr.Checkbox(label="Make new repo private?", value=True),
        gr.Checkbox(label="Overwrite existing repo?", value=True),
        gr.Checkbox(label="Create subdirectories automatically?", value=True),
    ],
    outputs=[
        gr.Markdown(label="output"),
        gr.Image(show_label=False),
    ],
    title="Duplicate your repo!",
    description="Duplicate a Hugging Face repository! This Space is a an experimental demo.",
    allow_flagging="never",
    live=False
)

def swap_visibilty(profile: gr.OAuthProfile | None):
    return gr.update(elem_classes=["main_ui_logged_in"]) if profile else gr.update(elem_classes=["main_ui_logged_out"])
        
css = '''
.main_ui_logged_out{opacity: 0.3; pointer-events: none}
'''
with gr.Blocks(css=css) as demo:
    gr.LoginButton()
    with gr.Column(elem_classes="main_ui_logged_out") as main_ui:
        interface.render()
    demo.load(fn=swap_visibilty, outputs=main_ui)
    
demo.queue()
demo.launch()