import os
import gradio as gr
from groq import Groq

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

def normalize_chat_history(chat_history):
    """
    Converts chat_history into list of tuples (user, assistant)
    Works with both tuple and dict formats returned by Gradio.
    """
    normalized = []
    for item in chat_history:
        if isinstance(item, dict):
            # If it's a dict, convert accordingly
            if item.get("role") == "user":
                normalized.append((item.get("content"), ""))
            elif item.get("role") == "assistant":
                # merge assistant into last user entry
                if normalized:
                    last_user, _ = normalized[-1]
                    normalized[-1] = (last_user, item.get("content"))
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            normalized.append((item[0], item[1]))
    return normalized

def chat_with_groq(user_input, chat_history):
    if not user_input.strip():
        return chat_history, ""

    chat_history = normalize_chat_history(chat_history)

    messages = []
    for u, a in chat_history:
        if u:
            messages.append({"role": "user", "content": u})
        if a:
            messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": user_input})

    try:
        chat_completion = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
        )

        assistant_response = chat_completion.choices[0].message.content
        chat_history.append((user_input, assistant_response))

    except Exception as e:
        chat_history.append((user_input, f"❌ Error: {e}"))

    return chat_history, ""

# UI Styling
custom_css = """
.gradio-container {
    max-width: 1000px;
    margin: auto;
    padding: 16px;
}

body {
    background: #f8f9fa;
    color: #1f2937;
    font-family: 'Inter', sans-serif;
}

h1 {
    color: #0b4f99;
    font-weight: bold;
    font-size: 2rem;
    text-align: center;
    margin-bottom: 12px;
}

#chatbot {
    background: #ffffff;
    border: 1px solid #d1d5db;
    border-radius: 12px;
    padding: 12px;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

.message.user {
    background: #0b4f99;
    color: #ffffff;
    border-radius: 16px;
    padding: 8px;
}

.message.bot {
    background: #e5e7eb;
    color: #1f2937;
    border-radius: 16px;
    padding: 8px;
}

textarea {
    width: 100%;
    border: 2px solid #6b7280;
    border-radius: 8px;
    padding: 8px;
    font-size: 1rem;
    color: #1f2937;
}

button {
    background: #0b4f99;
    color: #ffffff;
    font-weight: bold;
    border-radius: 8px;
    padding: 12px 20px;
}
button:hover, button:focus {
    background: #093c77;
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("<h1>🌐 Groq ChatBot</h1>")
    chatbot = gr.Chatbot(elem_id="chatbot", height=450)

    with gr.Row():
        msg = gr.Textbox(
            placeholder="💬 Type your message here...",
            show_label=False,
            scale=4,
        )
        send = gr.Button("Send", scale=1)

    clear = gr.Button("Clear Chat")

    send.click(chat_with_groq, [msg, chatbot], [chatbot, msg])
    msg.submit(chat_with_groq, [msg, chatbot], [chatbot, msg])
    clear.click(lambda: [], None, chatbot)

demo.launch()