import gradio as gr
import numpy as np
import config

from feifeilib.feifeichat import feifeichat
from feifeilib.feifeiimgtoimg import feifeiimgtoimg
from feifeilib.feifeitexttoimg import feifeitexttoimg
from feifeilib.feifeiflorence import feifeiflorence

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 4096

css = """
#col-container {
    width: auto;
    height: 750px;
}
"""


def create_ui():
    with gr.Blocks(css=css) as FeiFei:
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Tab("Generator"):
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        placeholder="Enter your prompt",
                        max_lines=12,
                        container=False,
                    )
                    run_button = gr.Button("Run")
                    result = gr.Image(
                        label="Result", show_label=False, interactive=False
                    )

                    with gr.Accordion("Advanced Settings", open=False):

                        seed = gr.Slider(
                            label="Seed",
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            value=0,
                        )

                        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                        with gr.Row():

                            width = gr.Slider(
                                label="Width",
                                minimum=256,
                                maximum=MAX_IMAGE_SIZE,
                                step=64,
                                value=896,
                            )

                            height = gr.Slider(
                                label="Height",
                                minimum=256,
                                maximum=MAX_IMAGE_SIZE,
                                step=64,
                                value=1152,
                            )

                        with gr.Row():
                            num_inference_steps = gr.Slider(
                                label="Number of inference steps",
                                minimum=1,
                                maximum=50,
                                step=1,
                                value=4,
                            )
                            guidancescale = gr.Slider(
                                label="Guidance scale",
                                minimum=0,
                                maximum=10,
                                step=0.1,
                                value=3.5,
                            )
                            num_strength = gr.Slider(
                                label="strength",
                                minimum=0,
                                maximum=2,
                                step=0.01,
                                value=0.35,
                            )
                with gr.Tab("Img2Img"):
                    img_in_result = gr.Image(label="Result", show_label=False)
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        placeholder="Enter your prompt",
                        max_lines=12,
                        container=False,
                    )
                    img_run_button = gr.Button("Img2Img")
                    img_out_result = gr.Image(
                        label="Result", show_label=False, interactive=False
                    )
                with gr.Tab("Styles"):
                    quality_select = gr.Checkbox(label="high quality")
                    sharpened_select = gr.Checkbox(label="Sharpened")
                    FooocusExpansion_select = gr.Checkbox(
                        label="FooocusExpansion", value=True
                    )
                    styles_name = [style["name"] for style in config.style_list]
                    styles_Radio = gr.Dropdown(
                        styles_name, label="Styles", multiselect=True
                    )
                with gr.Tab("Florence-2"):
                    with gr.Row():
                        with gr.Column():
                            output_text = gr.Textbox(label="Output Text", container=False)
                            florence_btn = gr.Button(value="Florence")
                        with gr.Column():
                            input_img = gr.Image(label="Input Picture")
            with gr.Column(scale=3, elem_id="col-container"):
                gr.ChatInterface(
                    feifeichat,
                    type="messages",
                    multimodal=True,
                    additional_inputs=[
                        gr.Checkbox(label="Feifei"),
                    ],
                )

        run_button.click(
            fn=feifeitexttoimg,  # Function to run for this button
            inputs=[
                prompt,
                quality_select,
                sharpened_select,
                styles_Radio,
                FooocusExpansion_select,
                seed,
                randomize_seed,
                width,
                height,
                num_inference_steps,
                guidancescale,
                num_strength,
            ],
            outputs=[result, seed],
        )

        img_run_button.click(
            fn=feifeiimgtoimg,  # Function to run for this button
            inputs=[img_in_result, prompt],
            outputs=[img_out_result],  # Output components for this function
        )

        florence_btn.click(
            fn=feifeiflorence,  # Function to run when the button is clicked
            inputs=[input_img],  # Input components for the function
            outputs=[output_text],  # Output component for the function
        )
    return FeiFei