import gradio as gr
from pathlib import Path
from huggingface_hub import list_repo_files, hf_hub_url
from collections import defaultdict
import requests

repo_id = 'nateraw/stable-diffusion-gallery'

def get_data(revision='main'):
    data = defaultdict(list)
    for file in list_repo_files(repo_id, repo_type='dataset', revision=revision):
        path = Path(file)
        if path.name == '.gitattributes':
            continue
        if path.suffix in ['.png', '.jpg', '.jpeg']:
            data[path.parent.name].append(file)

    prompts = []
    for key in sorted(data.keys()):
        prompt_cfg_url = hf_hub_url(repo_id=repo_id, filename=f"{key}/prompt_config.json", repo_type='dataset')
        prompt_cfg = requests.get(prompt_cfg_url).json()
        text = prompt_cfg.get('prompt')
        prompts.append(text)

    return data, prompts


def assembleHTML(data):
    # Borrowed mostly from:
    # https://huggingface.co./spaces/sd-concepts-library/stable-diffusion-conceptualizer
    html_gallery = ''
    html_gallery = html_gallery+'''
    <div class="flex gr-gap gr-form-gap row gap-4 w-full flex-wrap" id="main_row">
    '''
    for key, values in data.items():
        html_gallery = html_gallery+ f'''
        <div class="gr-block gr-box relative w-full overflow-hidden border-solid border border-gray-200 gr-panel">
            <div class="output-markdown gr-prose" style="max-width: 100%;">
            <h3>
                <a href="https://huggingface.co./datasets/nateraw/stable-diffusion-gallery/blob/main/{key}/prompt_config.json" target="_blank">
                <code>{ key }</code>
                </a>
            </h3>
            </div>
            <div id="gallery" class="gr-block gr-box relative w-full overflow-hidden border-solid border border-gray-200">
            <div class="wrap svelte-17ttdjv opacity-0"></div>
            <div class="absolute left-0 top-0 py-1 px-2 rounded-br-lg shadow-sm text-xs text-gray-500 flex items-center pointer-events-none bg-white z-20 border-b border-r border-gray-100 dark:bg-gray-900">
                <span class="mr-2 h-[12px] w-[12px] opacity-80">
                <svg xmlns="http://www.w3.org/2000/svg" width="100%" height="100%" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather feather-image">
                    <rect x="3" y="3" width="18" height="18" rx="2" ry="2"></rect>
                    <circle cx="8.5" cy="8.5" r="1.5"></circle>
                    <polyline points="21 15 16 10 5 21"></polyline>
                </svg>
                </span> Images
            </div>
            <div class="overflow-y-auto h-full p-2" style="position: relative;">
                <div class="grid gap-2 grid-cols-2 sm:grid-cols-2 md:grid-cols-2 lg:grid-cols-2 xl:grid-cols-2 2xl:grid-cols-2 svelte-1g9btlg pt-6">
        '''
        for image in values[:4]:
            html_gallery = html_gallery + f'''    
            <button class="gallery-item svelte-1g9btlg">
                <img alt="" loading="lazy" class="h-full w-full overflow-hidden object-contain" src="{hf_hub_url(repo_id, filename=image, repo_type='dataset')}">
            </button>
            '''
        html_gallery = html_gallery+'''
                    </div>
                    <iframe style="display: block; position: absolute; top: 0; left: 0; width: 100%; height: 100%; overflow: hidden; border: 0; opacity: 0; pointer-events: none; z-index: -1;" aria-hidden="true" tabindex="-1" src="about:blank"></iframe>
                </div>
                </div>
            </div>
            '''
    html_gallery = html_gallery+'''
    </div>
    '''
    return html_gallery


css = '''
.gradio-container {font-family: 'IBM Plex Sans', sans-serif}
#top_title{margin-bottom: .5em}
#top_title h2{margin-bottom: 0; text-align: center}
/*#main_row{flex-wrap: wrap; gap: 1em; max-height: 550px; overflow-y: scroll; flex-direction: row}*/
#component-3{height: 760px; overflow: auto}
#component-9{position: sticky;top: 0;align-self: flex-start;}
@media (min-width: 768px){#main_row > div{flex: 1 1 32%; margin-left: 0 !important}}
.gr-prose code::before, .gr-prose code::after {content: "" !important}
::-webkit-scrollbar {width: 10px}
::-webkit-scrollbar-track {background: #f1f1f1}
::-webkit-scrollbar-thumb {background: #888}
::-webkit-scrollbar-thumb:hover {background: #555}
.gr-button {white-space: nowrap}
.gr-button:focus {
  border-color: rgb(147 197 253 / var(--tw-border-opacity));
  outline: none;
  box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
  --tw-border-opacity: 1;
  --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
  --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
  --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
  --tw-ring-opacity: .5;
}
#prompt_input{flex: 1 3 auto; width: auto !important;}
#prompt_area{margin-bottom: .75em}
#prompt_area > div:first-child{flex: 1 3 auto}
.animate-spin {
    animation: spin 1s linear infinite;
}
@keyframes spin {
    from {
        transform: rotate(0deg);
    }
    to {
        transform: rotate(360deg);
    }
}
#share-btn-container {
    display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
}
#share-btn {
    all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;
}
#share-btn * {
    all: unset;
}
'''

data, prompts = get_data()

with gr.Blocks(css=css) as demo:
    with gr.Box().style(border=None):
        html = gr.HTML(assembleHTML(data))

demo.launch(debug=True)