import gradio as gr
from groq import Groq
import os

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

def generate_architecture(prompt, role, temperature):
    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[
            {"role": "system", "content": f"You are an expert {role}."},
            {"role": "user", "content": prompt}
        ],
        temperature=temperature
    )

    return completion.choices[0].message.content


with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown("""
    # 🚀 AI DevOps & MLOps Architect
    Design Cloud, Terraform, CI/CD & ML Architectures Instantly
    """)

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Ask your architecture question",
                placeholder="Example: Design scalable MLflow on AWS using Terraform...",
                lines=5
            )

        with gr.Column(scale=1):
            role = gr.Dropdown(
                choices=[
                    "DevOps Architect",
                    "MLOps Engineer",
                    "Cloud Security Engineer",
                    "Platform Engineer"
                ],
                value="DevOps Architect",
                label="Expert Mode"
            )

            temperature = gr.Slider(0, 1, value=0.7, label="Creativity")

            generate_btn = gr.Button("Generate 🚀")

    output = gr.Markdown(label="Generated Architecture")

    generate_btn.click(
        generate_architecture,
        inputs=[prompt, role, temperature],
        outputs=output
    )

demo.launch()