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_KEY")

# Initialize Groq client
client = Groq(api_key=api_key) if api_key else None

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}"


# 🎨 Define a custom theme
custom_theme = gr.themes.Base(
    primary_hue="blue",
    secondary_hue="violet",
    neutral_hue="slate",
    font=["Source Sans Pro", "Arial", "sans-serif"],
).set(
    body_background_fill="#0f172a",   # dark navy background
    body_text_color="#f1f5f9",        # light text
    button_primary_background_fill="#6366f1",  # indigo buttons
    button_primary_text_color="#ffffff",
    input_background_fill="#1e293b",  # dark input box
    input_border_color="#6366f1",
)

# Apply theme to ChatInterface
ui = gr.ChatInterface(
    fn=chat_groq,
    title="Affan ChatBot"
)

# Apply theme here
ui.launch(theme=custom_theme)