#!/usr/bin/env python

from __future__ import annotations

import enum

import gradio as gr
from huggingface_hub import HfApi

from inference import InferencePipeline
from utils import find_exp_dirs

SAMPLE_MODEL_IDS = [
    'koala2/dreambooth-dog-v2',
    'lambdalabs/dreambooth-avatar',
]


class ModelSource(enum.Enum):
    SAMPLE = 'Sample'
    HUB_LIB = 'Hub (dreambooth-library)'
    LOCAL = 'Local'


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

    @staticmethod
    def load_sample_model_list():
        return gr.update(choices=SAMPLE_MODEL_IDS, value=SAMPLE_MODEL_IDS[0])
    
    def load_hub_model_list(self) -> dict:
        api = HfApi(token=self.hf_token)
        choices = [
            info.modelId for info in api.list_models(author='dreambooth-library')
        ]
        return gr.update(choices=choices,
                         value=choices[0] if choices else None)

    @staticmethod
    def load_local_model_list() -> dict:
        choices = find_exp_dirs()
        return gr.update(choices=choices,
                         value=choices[0] if choices else None)

    def reload_model_list(self, model_source: str) -> dict:
        if model_source == ModelSource.SAMPLE.value:
            return self.load_sample_model_list()
        elif model_source == ModelSource.HUB_LIB.value:
            return self.load_hub_model_list()
        elif model_source == ModelSource.LOCAL.value:
            return self.load_local_model_list()
        else:
            raise ValueError

    def load_model_info(self, model_id: str) -> tuple[str, str]:
        try:
            card = InferencePipeline.get_model_card(model_id, self.hf_token)
        except Exception:
            return ''
        instance_prompt = getattr(card.data, 'instance_prompt', '')
        return instance_prompt

    def reload_model_list_and_update_model_info(
            self, model_source: str
        ) -> tuple[dict, str, str]:
        model_list_update = self.reload_model_list(model_source)
        model_list = model_list_update['choices']
        model_info = self.load_model_info(model_list[0] if model_list else '')
        return model_list_update, *model_info


def create_inference_demo(pipe: InferencePipeline,
                          hf_token: str | None = None) -> gr.Blocks:
    app = InferenceUtil(hf_token)

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                with gr.Box():
                    model_source = gr.Radio(
                        label='Model Source',
                        choices=[_.value for _ in ModelSource],
                        value=ModelSource.SAMPLE.value)
                    reload_button = gr.Button('Reload Model List')
                    model_id = gr.Dropdown(label='Model ID',
                                                choices=SAMPLE_MODEL_IDS,
                                                value=SAMPLE_MODEL_IDS[0])
                    with gr.Accordion(
                            label=
                            'Model info (Base model and instance prompt used for training)',
                            open=False):
                        with gr.Row():
                            instance_prompt_used_for_training = gr.Text(
                                label='Instance prompt', interactive=False)
                prompt = gr.Textbox(
                    label='Prompt',
                    max_lines=1,
                    placeholder='Example: "A picture of a {}dog in a bucket"'
                )
                seed = gr.Slider(label='Seed',
                                 minimum=0,
                                 maximum=100000,
                                 step=1,
                                 value=0)
                with gr.Accordion('Other Parameters', open=False):
                    num_steps = gr.Slider(label='Number of Steps',
                                          minimum=0,
                                          maximum=100,
                                          step=1,
                                          value=25)
                    guidance_scale = gr.Slider(label='CFG Scale',
                                               minimum=0,
                                               maximum=50,
                                               step=0.1,
                                               value=7.5)

                run_button = gr.Button('Generate')

                gr.Markdown('''
                - After training, you can press "Reload Model List" button to load your trained model names.
                ''')
            with gr.Column():
                result = gr.Image(label='Result')

        model_source.change(
            fn=app.reload_model_list_and_update_model_info,
            inputs=model_source,
            outputs=[
                model_id,
                instance_prompt_used_for_training,
            ])
        reload_button.click(
            fn=app.reload_model_list_and_update_model_info,
            inputs=model_source,
            outputs=[
                model_id,
                instance_prompt_used_for_training,
            ])
        model_id.change(fn=app.load_model_info,
                             inputs=model_id,
                             outputs=[
                                 instance_prompt_used_for_training,
                             ])
        inputs = [
            model_id,
            prompt,
            seed,
            num_steps,
            guidance_scale,
        ]
        prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
        run_button.click(fn=pipe.run, inputs=inputs, outputs=result)
    return demo


if __name__ == '__main__':
    import os

    hf_token = os.getenv('HF_TOKEN')
    pipe = InferencePipeline(hf_token)
    demo = create_inference_demo(pipe, hf_token)
    demo.queue(max_size=10).launch(share=False)