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
from feifeilib.feifeiflorencebase import process_image

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


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


def create_ui():
    with gr.Blocks(css=css) as FeiFei:
        with gr.Row():
            with gr.Column(scale=3):
                with gr.Tab("FeiFei"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            prompt = gr.Text(
                                label="Prompt",
                                show_label=False,
                                placeholder="Enter your prompt",
                                value="real girl in real life, ",
                                max_lines=12,
                                container=False,
                            )
                            feifei_button = gr.Button("FeiFei")
                            quality_select = gr.Checkbox(label="high quality")
                            sharpened_select = gr.Checkbox(label="Sharpened")
                            FooocusExpansion_select = gr.Checkbox(
                                label="Expansion", value=True)
                            styles_name = [
                                style["name"] for style in config.style_list
                            ]
                            styles_Radio = gr.Dropdown(styles_name,
                                                       label="Styles",
                                                       multiselect=False,
                                                       value="Photographic")
                            nsfw_select = gr.Checkbox(label="NSFW") 
                            nsfw_slider = gr.Slider(
                                label="NSFW",
                                minimum=0,
                                maximum=2,
                                step=0.05,
                                value=0.45,
                            )
                            out_prompt = gr.Text(
                                label="Prompt",
                                show_label=False,
                                max_lines=12,
                                placeholder="this photo prompt",
                                value="",
                                container=False,
                            )
                            with gr.Accordion("More",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)
    
                                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,
                                )
    
                                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.001,
                                    value=0.035,
                                )
                                
                                num_feifei = gr.Slider(
                                    label="FeiFei",
                                    minimum=0,
                                    maximum=2,
                                    step=0.05,
                                    value=0.45,
                                )
                            
                        with gr.Column(scale=2):
                            result = gr.Image(label="Result",
                                              show_label=False,
                                              interactive=False,
                                              height=940)

                with gr.Tab("GenPrompt"):

                    input_img = gr.Image(label="Input Picture",
                                         show_label=False,
                                         height=320,
                                         type="filepath")

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

                    output_text = gr.Textbox(label="Output Text",
                                             show_label=False,
                                             container=False)
                    
                with gr.Tab(label="Florence-2"):
                    with gr.Row():
                        florence_input_img = gr.Image(label="Input Picture",height=320,type="filepath")
                    with gr.Row():
                        florence_submit_btn = gr.Button(value="GenPrompt")
                    with gr.Row():
                        florence_output_text = gr.Textbox(label="Flux Prompt",
                                                          show_label=False,
                                                          container=False)
                        
            with gr.Column(scale=1, elem_id="col-container"):
                gr.ChatInterface(
                    feifeichat,
                    type="messages",
                    multimodal=True,
                    additional_inputs=[
                        gr.Checkbox(label="Feifei", value=True),
                        gr.Dropdown(
                            ["meta-llama/Llama-3.3-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="mistralai/Mistral-Nemo-Instruct-2411",
                            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,
                nsfw_select, 
                nsfw_slider,
                seed,
                randomize_seed,
                width,
                height,
                num_inference_steps,
                guidancescale,
                num_strength,
                num_feifei,
                
            ],
            outputs=[result, out_prompt],
        )

        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
        )

        florence_submit_btn.click(process_image, [florence_input_img], [florence_output_text])
    return FeiFei