import gradio as gr
from src.utils import add_to_prompt
from src.rep_api import generate_image_replicate, model_dict
import os

# The dictionary data
prompt_dict = {
    "Character": ["Asian girl with black hair", "A man with blond hair", "A Cat girl anime character with purple hair", "A Green Alien with big black eyes"],
    "Clothes": ["Wearing a blue jacket", "Wearing a black business suit", "Wearing a purple jumpsuit", "Wearing shorts and a white T-shirt"],
    "Pose": ["Close up portrait", "Standing doing a peace sign", "Folding arms", "holding a phone"],
    "Style": ["Simple white background", "Fashion runway", "Inside a business conference", "Inside a spaceship"]
}


def create_gen_tab():
    with gr.TabItem("Image Generator"):
        with gr.Row():
            with gr.Column():
                inp = gr.Textbox(label="Prompt")
                
                replicate_api_key = gr.Textbox(
                    label="Replicate API Key",
                    info="API key for Replicate",
                    value=os.environ.get("REPLICATE_API_TOKEN", ""),
                    type="password",
                )
                btn = gr.Button("Generate")
            with gr.Column():
                aspect_ratio = gr.Dropdown(list(["1:1","16:9","9:16","5:3"]),value="1:1", label="Aspect Ratio", info="Aspect Ratio")
                style_mode = gr.Dropdown(list(model_dict.keys()),label="Style lore",value="Base",info="style model")
                api_path = gr.Textbox(label="Other_style",info="Style lora safytensor route goes here or replicate link goes here",value=None)

        with gr.Accordion("Prompt Support", open=False):
            for key, values in prompt_dict.items():
                with gr.Row():
                    gr.Button(key,interactive=False)
                    for value in values:
                        gr.Button(value).click(add_to_prompt, inputs=[inp, gr.Textbox(value,visible=False)], outputs=inp)

        with gr.Row():
            gen_out = gr.Image(label="Generated Image",type="filepath")
        with gr.Accordion("Gallery", open=False):
            gen_gallery = gr.Gallery(label="Generated Images Gallery", type="filepath")

        btn.click(generate_image_replicate, inputs=[inp,style_mode,api_path,aspect_ratio,gen_gallery], outputs=[gen_out,gen_gallery])