import gradio as gr
from stkey_gr import stkey_gr, KEYS_FILES, DEFAULT_KEYS_FILE, upload_keys_file

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: right; margin: 20px; }
"""

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("# Safetensors file Key Checker", elem_classes="title")
        with gr.Group():
            dl_url = gr.Textbox(label="Download URL(s)", placeholder="https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors\n...", value="", lines=2, max_lines=255)
            with gr.Row():
                is_validate = gr.Checkbox(label="Validate safetensors file keys", value=True)
                rfile = gr.Dropdown(label="Reference file for validation", choices=KEYS_FILES, value=DEFAULT_KEYS_FILE)
            with gr.Accordion("Advanced", open=False):
                upload_rfile = gr.UploadButton(label="Upload reference file", file_types=["text"], file_count="single")
                with gr.Row():
                    with gr.Column():
                        hf_token = gr.Textbox(label="Your HF read token (Optional)", 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")
                    with gr.Column():
                        civitai_key = gr.Textbox(label="Your Civitai Key (Optional)", 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")
        run_button = gr.Button(value="Check", variant="primary")
        with gr.Group():
            uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=[]) # hidden
            urls_md = gr.Markdown("<br><br>", elem_classes="result", visible=False)
            out_files = gr.Files(label="Output", interactive=False, value=[])
            with gr.Row():
                missing = gr.JSON(value=[], label="Missing keys")
                added = gr.JSON(value=[], label="Added keys")
            with gr.Row():
                keys = gr.JSON(value=[], label="All keys")
                metadata = gr.JSON(value={}, label="Metadata")
            clear_button = gr.Button(value="Clear Output", variant="secondary")
    gr.DuplicateButton(value="Duplicate Space")

    gr.on(
        triggers=[run_button.click],
        fn=stkey_gr,
        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, out_files, is_validate, rfile],
        outputs=[uploaded_urls, out_files, urls_md, metadata, keys, missing, added],
    )
    clear_button.click(lambda: ([], [], "<br><br>", {}, [], [], []), None, [uploaded_urls, out_files, urls_md, metadata, keys, missing, added], queue=False, show_api=False)
    upload_rfile.upload(upload_keys_file, [upload_rfile], [rfile], queue=False, show_api=False)

demo.queue()
demo.launch()