import gradio as gr
import numpy as np
import config

from feifeilib.feifeichat import feifeichat
from feifeilib.feifeitexttoimg import feifeitexttoimg
from feifeilib.feifeiflorence import feifeiflorence
from feifeilib.feifeifluxapi import feifeifluxapi

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

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("FeiFei"):
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        placeholder="Enter your prompt",
                        max_lines=12,
                        container=False,
                    )
                    feifei_button = gr.Button("FeiFei")
                    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=512,
                                maximum=MAX_IMAGE_SIZE,
                                step=64,
                                value=1088,
                            )

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

                        with gr.Row():
                            num_inference_steps = gr.Slider(
                                label="Number of inference steps",
                                minimum=1,
                                maximum=50,
                                step=1,
                                value=6,
                            )
                            guidancescale = gr.Slider(
                                label="Guidance scale",
                                minimum=0,
                                maximum=10,
                                step=0.1,
                                value=0,
                            )
                            num_strength = gr.Slider(
                                label="strength",
                                minimum=0,
                                maximum=2,
                                step=0.001,
                                value=0.035,
                            )

                with gr.Tab("Styles"):
                    quality_select = gr.Checkbox(label="high quality")
                    sharpened_select = gr.Checkbox(label="Sharpened")
                    FooocusExpansion_select = gr.Checkbox(
                        label="FooocusExpansion")
                    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"):

                    input_img = gr.Image(label="Input Picture",
                                         show_label=False)

                    florence_btn = gr.Button(value="Florence")

                    output_text = gr.Textbox(label="Output Text",
                                             max_lines=12,
                                             show_label=False,
                                             container=False)
                    output_size = gr.Textbox(label="Output size",
                                             max_lines=2,
                                             show_label=False,
                                             container=False)
                    output_img = gr.Image(label="Input Picture",
                                          interactive=False,
                                          show_label=False)
                with gr.Tab("Flux"):
                    flux_prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        placeholder="Enter your prompt",
                        max_lines=12,
                        container=False,
                    )
                    flux_button = gr.Button("Flux")
                    flux_result = gr.Image(label="Result",
                                      show_label=False,
                                      interactive=False)

                    with gr.Accordion("Advanced Settings", open=False):
                        flux_seed = gr.Slider(
                            label="Seed",
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            value=0,
                        )

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

                        with gr.Row():
                            flux_width = gr.Slider(
                                label="Width",
                                minimum=512,
                                maximum=MAX_IMAGE_SIZE,
                                step=64,
                                value=896,
                            )

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

                        with gr.Row():
                            flux_num_inference_steps = gr.Slider(
                                label="Number of inference steps",
                                minimum=1,
                                maximum=50,
                                step=1,
                                value=4,
                            )
                            flux_guidancescale = gr.Slider(
                                label="Guidance scale",
                                minimum=0,
                                maximum=10,
                                step=0.1,
                                value=3.5,
                            )
                            flux_num_strength = gr.Slider(
                                label="strength",
                                minimum=0,
                                maximum=2,
                                step=0.001,
                                value=0.035,
                            )
            with gr.Column(scale=3, elem_id="col-container"):
                gr.ChatInterface(
                    feifeichat,
                    type="messages",
                    multimodal=True,
                    additional_inputs=[
                        gr.Checkbox(label="Feifei",value=True),
                        gr.Dropdown(
                            ["meta-llama/Meta-Llama-3.1-70B-Instruct", 
                             "CohereForAI/c4ai-command-r-plus-08-2024",
                             "Qwen/Qwen2.5-72B-Instruct",
                             "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
                             "NousResearch/Hermes-3-Llama-3.1-8B",
                             "mistralai/Mistral-Nemo-Instruct-2411",
                             "microsoft/Phi-3.5-mini-instruct"],
                            value="CohereForAI/c4ai-command-r-plus-08-2024",
                            show_label=False,
                            container=False),
                        gr.Radio(
                            ["pixtral","Vsiion"],
                            value="pixtral",
                            show_label=False,
                            container=False)
                    ],
                )

        feifei_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],
        )
        flux_button.click(
            fn=feifeifluxapi,  # Function to run for this button
            inputs=[
                flux_prompt,
                flux_height,
                flux_width,
                flux_guidancescale
            ],
            outputs=[flux_result],
        )
        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_size,
                     output_img],  # Output component for the function
        )
    return FeiFei