#!/usr/bin/env python

from __future__ import annotations

import os
import random
from typing import Tuple, Optional

import gradio as gr
from huggingface_hub import HfApi

from inf import InferencePipeline


SAMPLE_MODEL_IDS = [
    'lora-library/B-LoRA-teddybear',
    'lora-library/B-LoRA-bull',
    'lora-library/B-LoRA-wolf_plushie',
    'lora-library/B-LoRA-pen_sketch',
    'lora-library/B-LoRA-cartoon_line',
    'lora-library/B-LoRA-multi-dog2',
]
css = """
body {
    font-size: 30px;
}

#title {
    text-align: center;
}

#title h1 {
    font-size: 250%;
}

.lora-title {
    text-align: center;
    border-radius: 10px;
    background: linear-gradient(90deg, #1CB5E0 0%, #000851 100%);
}

.gr-image {
    width: 512px;
    height: 512px;
    object-fit: contain;
    margin: auto;
}

.lora-column {
    display: flex;
    flex-direction: column;
    align-items: center; 
    justify-content: center;
    border: none;
    background: none;
}
.gr-row {
    align-items: center;
    justify-content: center;
    margin-top: 5px;
}
"""

def get_choices(hf_token):
    api = HfApi(token=hf_token)
    choices = [
        info.modelId for info in api.list_models(author='lora-library')
    ]
    models_list = ['None'] + SAMPLE_MODEL_IDS + choices
    return models_list


def get_image_from_card(card, model_id) -> Optional[str]:
    try:
        card_path = f"https://huggingface.co/{model_id}/resolve/main/"
        widget = card.data.get('widget')
        if widget is not None or len(widget) > 0:
            output = widget[0].get('output')
            if output is not None:
                url = output.get('url')
                if url is not None:
                    return card_path + url
        return None
    except Exception:
        return None


def demo_init():
    try:
        choices = get_choices(app.hf_token)
        content_blora = random.choice(SAMPLE_MODEL_IDS)
        style_blora = random.choice(SAMPLE_MODEL_IDS)
        content_blora_prompt, content_blora_image = app.load_model_info(content_blora)
        style_blora_prompt, style_blora_image = app.load_model_info(style_blora)

        content_lora_model_id = gr.update(choices=choices, value=content_blora)
        content_prompt = gr.update(value=content_blora_prompt)
        content_image = gr.update(value=content_blora_image)

        style_lora_model_id = gr.update(choices=choices, value=style_blora)
        style_prompt = gr.update(value=style_blora_prompt)
        style_image = gr.update(value=style_blora_image)

        prompt = gr.update(
            value=f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style')

        return content_lora_model_id, content_prompt, content_image, style_lora_model_id, style_prompt, style_image, prompt

    except Exception as e:
        raise type(e)(f'failed to demo_init, due to: {e}')


def toggle_column(is_checked):
    try:
        return 'None' if is_checked else random.choice(SAMPLE_MODEL_IDS)
    except Exception as e:
        raise type(e)(f'failed to toggle_column, due to: {e}')


class InferenceUtil:
    def __init__(self, hf_token: str | None):
        self.hf_token = hf_token

    def load_model_info(self, lora_model_id: str) -> Tuple[str, Optional[str]]:
        try:
            try:
                card = InferencePipeline.get_model_card(lora_model_id,
                                                        self.hf_token)
            except Exception:
                return '', None
            instance_prompt = getattr(card.data, 'instance_prompt', '')
            image_url = get_image_from_card(card, lora_model_id)
            return instance_prompt, image_url
        except Exception as e:
            raise type(e)(f'failed to load_model_info, due to: {e}')

    def update_model_info(self, model_source: str):
        try:
            if model_source == 'None':
                return '', None
            else:
                model_info = self.load_model_info(model_source)
                new_prompt, new_image = model_info[0], model_info[1]
            return new_prompt, new_image
        except Exception as e:
            raise type(e)(f'failed to update_model_info, due to: {e}')


hf_token = os.getenv('HF_TOKEN')
pipe = InferencePipeline(hf_token)
app = InferenceUtil(hf_token)


with gr.Blocks(css=css) as demo:
    title = gr.HTML(
        '''<h1>Implicit Style-Content Separation using B-LoRA</h1>
        <p>This is a demo for our <a href="https://arxiv.org/abs/2403.14572">paper</a>: <b>''Implicit Style-Content Separation using B-LoRA''</b>.
    <br>
    Project page and code is available <a href="https://b-lora.github.io/B-LoRA/">here</a>.</p>
        ''',
        elem_id="title"
      )
    with gr.Row(elem_classes="gr-row"):
        with gr.Column():
            with gr.Group(elem_classes="lora-column"):
                content_sub_title = gr.HTML('''<h2>Content B-LoRA</h2>''', elem_classes="lora-title")
                content_checkbox = gr.Checkbox(label='Use Content Only', value=False)
                content_lora_model_id = gr.Dropdown(label='Model ID', choices=[])
                content_prompt = gr.Text(label='Content instance prompt', interactive=False, max_lines=1)
                content_image = gr.Image(label='Content Image', elem_classes="gr-image")
        with gr.Column():
            with gr.Group(elem_classes="lora-column"):
                style_sub_title = gr.HTML('''<h2>Style B-LoRA</h2>''', elem_classes="lora-title")
                style_checkbox = gr.Checkbox(label='Use Style Only', value=False)
                style_lora_model_id = gr.Dropdown(label='Model ID', choices=[])
                style_prompt = gr.Text(label='Style instance prompt', interactive=False, max_lines=1)
                style_image = gr.Image(label='Style Image', elem_classes="gr-image")
    with gr.Row(elem_classes="gr-row"):
        with gr.Column():
            with gr.Group():
                prompt = gr.Textbox(
                    label='Prompt',
                    max_lines=1,
                    placeholder='Example: "A [c] in [s] style"'
                )
                result = gr.Image(label='Result')
                with gr.Accordion('Other Parameters', open=False, elem_classes="gr-accordion"):
                    content_alpha = gr.Slider(label='Content B-LoRA alpha',
                                              minimum=0,
                                              maximum=2,
                                              step=0.05,
                                              value=1)
                    style_alpha = gr.Slider(label='Style B-LoRA alpha',
                                            minimum=0,
                                            maximum=2,
                                            step=0.05,
                                            value=1)
                    seed = gr.Slider(label='Seed',
                                     minimum=0,
                                     maximum=100000,
                                     step=1,
                                     value=8888)
                    num_steps = gr.Slider(label='Number of Steps',
                                          minimum=0,
                                          maximum=100,
                                          step=1,
                                          value=50)
                    guidance_scale = gr.Slider(label='CFG Scale',
                                               minimum=0,
                                               maximum=50,
                                               step=0.1,
                                               value=7.5)
                run_button = gr.Button('Generate')
    demo.load(demo_init, inputs=[],
              outputs=[content_lora_model_id, content_prompt, content_image, style_lora_model_id, style_prompt,
                       style_image, prompt], queue=False, show_progress="hidden")
    content_lora_model_id.change(
        fn=app.update_model_info,
        inputs=content_lora_model_id,
        outputs=[
            content_prompt,
            content_image,
        ])
    style_lora_model_id.change(
        fn=app.update_model_info,
        inputs=style_lora_model_id,
        outputs=[
            style_prompt,
            style_image,
        ])
    style_prompt.change(
        fn=lambda content_blora_prompt,
                  style_blora_prompt: f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style' if style_blora_prompt else content_blora_prompt,
        inputs=[content_prompt, style_prompt],
        outputs=prompt,
    )
    content_prompt.change(
        fn=lambda content_blora_prompt,
                  style_blora_prompt: f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style' if content_blora_prompt else style_blora_prompt,
        inputs=[content_prompt, style_prompt],
        outputs=prompt,
    )
    content_checkbox.change(toggle_column, inputs=[content_checkbox],
                            outputs=[style_lora_model_id])
    style_checkbox.change(toggle_column, inputs=[style_checkbox],
                          outputs=[content_lora_model_id])
    inputs = [
        content_lora_model_id,
        style_lora_model_id,
        prompt,
        content_alpha,
        style_alpha,
        seed,
        num_steps,
        guidance_scale,
    ]
    prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
    run_button.click(fn=pipe.run, inputs=inputs, outputs=result)

demo.queue(max_size=10).launch(share=False)