import os
import gradio as gr
from groq import Groq
import torch
from diffusers import StableDiffusionPipeline

# ===== GROQ CLIENT =====
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# ===== LOAD STABLE DIFFUSION =====
model_id = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
)

pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# ===== PROMPT ENHANCER USING GROQ =====
def enhance_prompt(user_prompt):
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {
                "role": "system",
                "content": (
                    "You are an expert AI image prompt engineer. "
                    "Enhance the user prompt with visual details, lighting, style, realism, "
                    "camera angle, and artistic quality. "
                    "Return ONLY the enhanced prompt."
                )
            },
            {"role": "user", "content": user_prompt}
        ],
        temperature=0.7,
        max_tokens=150
    )

    return response.choices[0].message.content.strip()

# ===== IMAGE GENERATION =====
def generate_image(prompt):
    if not prompt.strip():
        return None, ""

    enhanced_prompt = enhance_prompt(prompt)

    image = pipe(
        enhanced_prompt,
        num_inference_steps=30,
        guidance_scale=7.5
    ).images[0]

    return image, enhanced_prompt

# ===== UI =====
with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Generator (Groq + SD)") as demo:
    gr.Markdown(
        """
        # 🎨 AI Image Generator  
        **Groq LLM + Stable Diffusion**  
        Type a prompt → Get a stunning image
        """
    )

    prompt_input = gr.Textbox(
        label="Image Prompt",
        placeholder="e.g. A futuristic city at sunset",
        lines=3
    )

    generate_btn = gr.Button("Generate Image 🚀")

    image_output = gr.Image(label="Generated Image")
    enhanced_prompt_output = gr.Textbox(
        label="Enhanced Prompt (Generated by Groq)",
        lines=4
    )

    generate_btn.click(
        generate_image,
        inputs=prompt_input,
        outputs=[image_output, enhanced_prompt_output]
    )

    gr.Markdown("### 🔹 Demo Prompts (Click & Generate)")

    demo_prompts = [
        "A cyberpunk city at night",
        "A realistic lion in the jungle",
        "A futuristic robot chef",
        "A fantasy castle in the clouds",
        "A Pakistani truck art style painting"
    ]

    for demo_text in demo_prompts:
        gr.Button(demo_text).click(
            generate_image,
            inputs=gr.State(demo_text),
            outputs=[image_output, enhanced_prompt_output]
        )

demo.launch()