# app.py 
from config.config import models, prompts, api_token  # Direct import
import gradio as gr
from src.img_gen import generate_image
import base64

# Gradio Interface
def gradio_interface():
    # LOAD CUSTOM CSS
    try:
        with open("config/layout.css", "r") as f:
            custom_css = f.read()
    except FileNotFoundError:
        print("Error: aaa.css not found!")
        custom_css = ""  # Or provide default CSS

    with gr.Blocks(css=custom_css) as demo:
        gr.Markdown("# ========== CtB AI Image Generator ==========")
        with gr.Row():
            # Set default values for dropdowns
            prompt_dropdown = gr.Dropdown(choices=[p["alias"] for p in prompts], label="Select Prompt", value=prompts[0]["alias"])
            team_dropdown = gr.Dropdown(choices=["Red", "Blue"], label="Select Team", value="Red")
            #model_dropdown = gr.Dropdown(choices=[m["alias"] for m in models], label="Select Model", value=models[0]["alias"])
        with gr.Row():
            # Add a text box for custom user input (max 200 characters)
            custom_prompt_input = gr.Textbox(label="Custom Prompt (Optional)", placeholder="Enter additional details (max 200 chars)...", max_lines=1, max_length=200)
        with gr.Row():
            generate_button = gr.Button("Generate Image")
        with gr.Row():
            output_image = gr.Image(elem_classes="output-image", label="Generated Image", show_label=False, scale=1, width="100%")
        with gr.Row():
            status_text = gr.Textbox(label="Status", placeholder="Waiting for input...", interactive=False)

        # Connect the button to the function
        generate_button.click(
            generate_image,
            inputs=[prompt_dropdown, 
                    team_dropdown, 
                    custom_prompt_input,
                    #model_dropdown,
                    ],
            outputs=[output_image, status_text]
        )
    return demo

# Create the demo instance
demo = gradio_interface()

# Only launch if running directly
if __name__ == "__main__":
        demo.queue().launch()