import os
import gradio as gr
from groq import Groq

api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    raise ValueError("GROQ_API_KEY not found in environment variables.")

client = Groq(api_key=api_key)
MODEL_NAME = "llama-3.1-8b-instant"

def normalize_history(history):
    """
    Convert Gradio history to (user, assistant) pairs.
    Handles both:
    - [["hi", "hello"], ["ok", "sure"]]
    - [{"role":"user","content":"hi"}, {"role":"assistant","content":"hello"}, ...]
    """
    normalized = []

    if not history:
        return normalized

    if isinstance(history[0], list):
        # Format: [["hi", "hello"], ["ok", "sure"]]
        return history

    # Format: [{"role": "...", "content": "..."}]
    temp_pair = []
    for msg in history:
        if msg["role"] == "user":
            temp_pair = [msg["content"], None]
        elif msg["role"] == "assistant":
            if temp_pair:
                temp_pair[1] = msg["content"]
                normalized.append(temp_pair)
                temp_pair = []

    return normalized


def chat_groq(message, history):
    try:
        history_pairs = normalize_history(history)

        system_prompt = "You are a helpful AI assistant."
        messages = [{"role": "system", "content": system_prompt}]

        # Add normalized history
        for user_msg, bot_msg in history_pairs:
            messages.append({"role": "user", "content": user_msg})
            if bot_msg:
                messages.append({"role": "assistant", "content": bot_msg})

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

        # Groq API
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=400,
            temperature=0.7,
        )

        return completion.choices[0].message.content

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


ui = gr.ChatInterface(
    fn=chat_groq,
    title="Sajid's GenAI App (Groq)",
    description="Simple Llama-3 chat app running on Groq.",
)

if __name__ == "__main__":
    ui.launch()