import os
import gradio as gr
from groq import Groq

# Read API key from environment (Hugging Face Spaces secret)
api_key = os.environ.get("GROQ_API")

if not api_key:
    raise ValueError("No Groq API key found. Please add it as a secret.")

# Initialize Groq client
client = Groq(api_key=api_key)

MODEL_NAME = "openai/gpt-oss-safeguard-20b"

def chat_groq(message, history):
    try:
        system_prompt = "You are a helpful AI assistant."
        messages = [{"role": "system", "content": system_prompt}]

        # Handle history safely
        for turn in history:
            if isinstance(turn, list) and len(turn) == 2:
                user_msg, bot_msg = turn
                messages.append({"role": "user", "content": user_msg})
                if bot_msg:
                    messages.append({"role": "assistant", "content": bot_msg})

        # Add latest user message
        messages.append({"role": "user", "content": message})

        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=1024,
            temperature=0.7
        )

        reply = completion.choices[0].message.content
        return reply

    except Exception as e:
        return f"Error: {e}"


ui = gr.ChatInterface(
    fn=chat_groq,
    title="Affan's GenAI App(Groq)"
)

ui.launch()