import gradio as gr
from civitai_to_hf import (search_civitai, download_civitai, select_civitai_item, add_civitai_item, get_civitai_tag,
                           CIVITAI_TYPE, CIVITAI_BASEMODEL, CIVITAI_SORT, CIVITAI_PERIOD)

css = """
.title { font-size: 3em; align-items: center; text-align: center; }
.info { align-items: center; text-align: center; }
.block.result { margin: 1em 0; padding: 1em; box-shadow: 0 0 3px 3px #664422, 0 0 3px 2px #664422 inset; border-radius: 6px; background: #665544; }
.desc [src$='#float'] { float: left; margin: 10px; }
"""

with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
    with gr.Column():
        gr.Markdown("# CivitAI to HF🤗 Downloader & Uploader", elem_classes="title")
        with gr.Accordion("Search Civitai", open=True):
            with gr.Row():
                search_civitai_type = gr.CheckboxGroup(label="Type", choices=CIVITAI_TYPE, value=["Checkpoint", "LORA"])
                search_civitai_basemodel = gr.CheckboxGroup(label="Base model", choices=CIVITAI_BASEMODEL, value=[])
            with gr.Row():
                search_civitai_sort = gr.Radio(label="Sort", choices=CIVITAI_SORT, value=CIVITAI_SORT[0])
                search_civitai_period = gr.Radio(label="Period", choices=CIVITAI_PERIOD, value="Month")
            with gr.Row():
                search_civitai_query = gr.Textbox(label="Query", placeholder="flux", lines=1)
                search_civitai_tag = gr.Dropdown(label="Tag", choices=get_civitai_tag(), value=get_civitai_tag()[0], allow_custom_value=True)
                search_civitai_user = gr.Textbox(label="Username", lines=1)
                search_civitai_submit = gr.Button("Search on Civitai")
            with gr.Row():
                search_civitai_desc = gr.Markdown(value="", visible=False, elem_classes="desc")
                search_civitai_json = gr.JSON(value={}, visible=False)
            search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value="", allow_custom_value=True, visible=False, multiselect=True)
            search_civitai_add = gr.Button("Add to download URLs")
        with gr.Group():
            dl_url = gr.Textbox(label="Download URL(s)", placeholder="https://civitai.com/api/download/models/28907\n...", value="", lines=3, max_lines=255)
            with gr.Column():
                civitai_key = gr.Textbox(label="Your Civitai Key", value="", max_lines=1)
                gr.Markdown("Your Civitai API key is available at [https://civitai.com/user/account](https://civitai.com/user/account).", elem_classes="info")
        with gr.Group():
            with gr.Row():
                with gr.Column():
                    hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
                    gr.Markdown("Your token is available at [hf.co/settings/tokens](https://huggingface.co/settings/tokens).", elem_classes="info")
                newrepo_id = gr.Textbox(label="Upload repo ID", placeholder="yourid/yourrepo", value="", max_lines=1)
            with gr.Row():
                newrepo_type = gr.Radio(label="Upload repo type", choices=["model", "dataset"], value="model")
                is_private = gr.Checkbox(label="Create private repo", value=True)
                is_info = gr.Checkbox(label="Upload Civitai information files", value=False)
        run_button = gr.Button(value="Download and Upload", variant="primary")
        uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
        urls_md = gr.Markdown("<br><br>", elem_classes="result")
    gr.DuplicateButton(value="Duplicate Space")

    gr.on(
        triggers=[run_button.click],
        fn=download_civitai,
        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, newrepo_id, newrepo_type, is_private, is_info],
        outputs=[uploaded_urls, urls_md],
    )
    gr.on(
        triggers=[search_civitai_submit.click, search_civitai_query.submit, search_civitai_user.submit],
        fn=search_civitai,
        inputs=[search_civitai_query, search_civitai_type, search_civitai_basemodel, search_civitai_sort, search_civitai_period, search_civitai_tag, search_civitai_user],
        outputs=[search_civitai_result, search_civitai_desc, search_civitai_submit, search_civitai_query],
        queue=True,
        show_api=False,
    )
    search_civitai_result.change(select_civitai_item, [search_civitai_result], [search_civitai_desc, search_civitai_json], queue=False, show_api=False)
    search_civitai_add.click(add_civitai_item, [search_civitai_result, dl_url], [dl_url], queue=False, show_api=False)

demo.queue()
demo.launch()