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_title, token, dest_namespace=None, 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"Copied from {collection.title} using https://huggingface.co/spaces/librarian-bots/collection_cloner."
    if dest_namespace == "username":
        dest_namespace = None
    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(css="style.css") 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 or URL",
            placeholder="e.g. username/collection-slug",
        )
        gr.Markdown("## Destination Collection info")

        dest_title = gr.Textbox(
            label="Destination Title",
        )
        dest_namespace = gr.Textbox(
            value="username",
            label="Destination Namespace (optional - defaults to your username))",
            interactive=True,
        )
        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_title,
            token,
            dest_namespace,
            private,
            overwrite,
        ],
        response,
    )

demo.launch()