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

# ----------------------------------------------------
# 1. GROQ CHATBOT  (uses HF Secret: GROQ_API_KEY)
# ----------------------------------------------------
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

def chat_with_groq(text):
    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[{"role": "user", "content": text}]
    )
    return completion.choices[0].message.content


# ----------------------------------------------------
# 2. CPU IMAGE GENERATOR (uses SD-Turbo, CPU friendly)
# ----------------------------------------------------
pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sd-turbo",
    torch_dtype=torch.float32
).to("cpu")

def generate_image(prompt):
    img = pipe(prompt, num_inference_steps=6).images[0]
    return img


# ----------------------------------------------------
# 3. GRADIO MULTI-TAB UI
# ----------------------------------------------------
with gr.Blocks() as app:
    gr.Markdown("## 🤖 Groq Chatbot + 🎨 Image Generator (CPU Version)")

    with gr.Tab("💬 Chatbot"):
        in_txt = gr.Textbox(label="Enter text", lines=4)
        out_txt = gr.Textbox(label="Response", lines=10)
        btn_chat = gr.Button("Send")
        btn_chat.click(chat_with_groq, inputs=in_txt, outputs=out_txt)

    with gr.Tab("🎨 Image Generator"):
        prompt = gr.Textbox(label="Enter a prompt")
        img_out = gr.Image(label="Generated Image")
        btn_img = gr.Button("Generate Image")
        btn_img.click(generate_image, inputs=prompt, outputs=img_out)

app.launch()