import os
from huggingface_hub import InferenceClient
import gradio as gr

from utils import get_default_hyperparameters  # Import the utility function


hf_token = os.getenv("HF_TOKEN")

headers = {
    "x-wait-for-model": "true",
    "x-use-cache": "true", # Use past generations
}

imagegen_models_list = [
    "black-forest-labs/FLUX.1-schnell",
    "black-forest-labs/FLUX.1-dev",
    "strangerzonehf/Flux-Midjourney-Mix2-LoRA",
    "stabilityai/stable-diffusion-3.5-large",
    "stabilityai/stable-diffusion-xl-base-1.0",
    "stable-diffusion-v1-5/stable-diffusion-v1-5"
]

promptgen_models_list = [
    "meta-llama/Llama-3.2-11B-Vision-Instruct",
    "mistralai/Mistral-Nemo-Instruct-2407"
]

# Function to generate image
def generate_image(model, prompt, guidance, width, height, num_inference_steps, seed):
    print(f"""Generating image with following parameters: 
          Model: {model}, Prompt: {prompt}, Guidance Scale: {guidance}, Width: {width}, Height: {height},""")
    client = InferenceClient(model=model, headers=headers, token=hf_token)
    image = client.text_to_image(
        model=model,
        prompt=prompt,
        guidance_scale=guidance,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        seed=seed
    )
    return image

# Function to update hyperparameters dynamically
def update_hyperparameters(model_name):
    default_params = get_default_hyperparameters(model_name)
    return (default_params['guidance_scale'], 
            default_params['width'], 
            default_params['height'], 
            default_params['num_inference_steps'], 
            default_params['seed'])

# Function to expand the idea using the selected prompt generation model
def expand_idea(promptgen_model, idea_text):
    print(f"Expanding idea with model: {promptgen_model}")
    client = InferenceClient(model=promptgen_model, headers=headers, token=hf_token)
    response = client.chat_completion(
        messages=[
            {
                'role': 'user',
                'content': f'For the given idea, generate a text prompt to generate an image from a text to image generator. Be creative and include both subject and style prompts into one. Do not, explain your decisions. Idea: {idea_text}',
            },
        ],
        max_tokens=100,
        temperature=1.1,  # Set temperature higher for dynamic responses
        top_p=0.9,
    )
    expanded_prompt = response.choices[0].message.content  # Assuming this is how the prompt is expanded
    return expanded_prompt

# Interface for generating images and expanding idea into a prompt
def run_interface():
    with gr.Blocks() as iface:
        with gr.Row():
            # Image generation controls
            with gr.Column(scale=1):
                model_dropdown = gr.Dropdown(choices=imagegen_models_list, label="Image Model", value=imagegen_models_list[0])
                prompt_textbox = gr.Textbox(label="Prompt", lines=5, value="Astronaut floating in space")
                
                # Initial default values based on the first model
                default_params = get_default_hyperparameters(model_dropdown.value)

                guidance_slider = gr.Slider(0, 10, step=0.1, label="Guidance Scale", value=default_params['guidance_scale'])
                width_slider = gr.Slider(256, 2048, step=32, label="Width", value=default_params['width'])
                height_slider = gr.Slider(256, 2048, step=32, label="Height", value=default_params['height'])
                steps_slider = gr.Slider(1, 100, step=1, label="Number of Inference Steps", value=default_params['num_inference_steps'])
                seed_number = gr.Number(label="Seed", value=default_params['seed'])

                # Update sliders based on model selection
                model_dropdown.change(
                    fn=update_hyperparameters,
                    inputs=model_dropdown,
                    outputs=[guidance_slider, width_slider, height_slider, steps_slider, seed_number]
                )

                generate_button = gr.Button("Generate Image")
                output_image = gr.Image(type="pil", format="png")

                generate_button.click(
                    fn=generate_image,
                    inputs=[model_dropdown, prompt_textbox, guidance_slider, width_slider, height_slider, steps_slider, seed_number],
                    outputs=output_image
                )

            # Prompt expansion section in a collapsible panel
            with gr.Column(scale=1):
                gr.Markdown("**Expand an Idea into Prompt**")
                with gr.Accordion("Your Image Idea", open=False):
                    promptgen_dropdown = gr.Dropdown(choices=promptgen_models_list, label="Prompt Generation Model", value=promptgen_models_list[0])
                    idea_textbox = gr.Textbox(label="Enter your idea/sketch (rough text)", lines=5, placeholder="Describe your idea here...")
                    expand_button = gr.Button("Expand to Prompt")

                    # Expand idea and directly output to the main prompt textbox
                    expand_button.click(
                        fn=expand_idea,
                        inputs=[promptgen_dropdown, idea_textbox],
                        outputs=prompt_textbox  # Updating the prompt_textbox directly
                    )
    
    iface.launch()

run_interface()