from huggingface_hub import HfApi
import gradio as gr
import re


def extract_slug(url):
    pattern = r"https://huggingface\.co/collections/(.*)"
    return match.group(1) if (match := re.search(pattern, url)) else None


def clone_collection(
    source_slug, dest_namespace, dest_title, token, private=False, exist_ok=False
):
    api = HfApi(token=token)
    source_slug = source_slug.strip()
    # check if formatted as url
    if source_slug.startswith("https://huggingface.co./collections/"):
        source_slug = extract_slug(source_slug)
    collection = api.get_collection(source_slug)
    if not collection:
        raise gr.Error(
            f"Collection {source_slug} does not exist or you do not have access to it."
        )
    description = f"Copy of {collection.title}"

    new_collection = api.create_collection(
        dest_title,
        namespace=dest_namespace,
        exists_ok=exist_ok,
        private=private,
        description=description,
        token=token,
    )
    for item in collection.items:
        api.add_collection_item(
            new_collection.slug, item.item_id, item_type=item.item_type
        )
    return f"[Collection]({collection.url}) has been cloned into [new_collection.slug]({new_collection.url})"


title = (
    """<h1 style='text-align: center;'> &#129516; Collection Cloner &#129516;</h1>"""
)


with gr.Blocks() as demo:
    gr.HTML(title)
    gr.HTML(
        """<p style='text-align: center;'>
This space allows you to clone a <a href="https://huggingface.co./docs/hub/collections">Collection</a> from the Hugging Face Hub into your own namespace.
You can then edit the collection to your liking!</p>"""
    )
    gr.Markdown("## Authentication")
    gr.Markdown(
        "Token is required to create a new collection and clone private collections. You can get your token from your [profile page](https://huggingface.co./settings/token)."
    )
    with gr.Row():
        token = gr.Textbox(
            label="Token",
            type="password",
        )
    with gr.Column():
        gr.Markdown("## Source Collection")
        source_slug = gr.Textbox(
            label="Source Collection Slug", placeholder="e.g. username/collection-slug"
        )
        gr.Markdown("## Destination Collection info")
        dest_namespace = gr.Textbox(
            label="Destination Namespace",
            interactive=True,
        )
        dest_title = gr.Textbox(
            label="Destination Title",
        )
        with gr.Row():
            private = gr.Checkbox(
                False,
                label="Make new collection private?",
            )
            overwrite = gr.Checkbox(
                False,
                label="Overwrite any collection with same slug as the destination?",
            )
    submit_btn = gr.Button("Clone Collection")
    response = gr.Markdown()
    submit_btn.click(
        clone_collection,
        [
            source_slug,
            dest_namespace,
            dest_title,
            token,
            private,
            overwrite,
        ],
        response,
    )

demo.launch()