import gradio as gr
import os
import spaces 
from groq import Groq

# Initialize the Groq client via Hugging Face Secrets
api_key = os.environ.get("GROQ_API_KEY")
if not api_key:
    raise ValueError("GROQ_API_KEY environment variable not found. Please add it as a Secret in your Space Settings.")

client = Groq(api_key=api_key)

# Custom CSS for modern styling
custom_css = """
.gradio-container { background-color: #0b0f19; font-family: 'Inter', sans-serif; }
#title-header { text-align: center; margin-bottom: 20px; }
#title-header h1 { color: #38bdf8; font-weight: 800; font-size: 2.2rem; }
.sidebar-panel { background: #111827 !important; border: 1px solid #1f2937 !important; border-radius: 12px !important; }
.chat-window { border: 1px solid #1f2937 !important; border-radius: 12px !important; background: #111827 !important; }
"""

@spaces.GPU
def chat_stream(message, history, model, system_prompt, temperature, max_tokens):
    """
    Handles streaming responses using the standard tuple-based chat history format.
    history layout: [[user_msg1, bot_msg1], [user_msg2, bot_msg2], ...]
    """
    if not message.strip():
        yield history
        return

    # 1. Initialize message list with system prompt for Groq API
    api_messages = [{"role": "system", "content": system_prompt}]

    # 2. Append existing conversation history seamlessly 
    for user_msg, bot_msg in history:
        if user_msg:
            api_messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            api_messages.append({"role": "assistant", "content": bot_msg})

    # 3. Append the newest user message to the API list
    api_messages.append({"role": "user", "content": message})
    
    # 4. Update the Gradio UI history with an empty slot for the bot's upcoming message
    history.append([message, ""])
    yield history

    # 5. Stream from Groq
    try:
        stream = client.chat.completions.create(
            model=model,
            messages=api_messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
        )

        partial_response = ""
        for chunk in stream:
            if chunk.choices[0].delta.content:
                partial_response += chunk.choices[0].delta.content
                # Update the very last bot message slot in history
                history[-1][1] = partial_response
                yield history

    except Exception as e:
        history[-1][1] = f"⚠️ Error connecting to Groq: {str(e)}"
        yield history

# Build the layout manually
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Personal AI ChatBot", elem_id="title-header")
    gr.Markdown("A fully customizable, hyper-fast LLM workspace.")

    with gr.Row():
        # --- LEFT COLUMN: CONTROL PANEL ---
        with gr.Column(scale=1, elem_classes="sidebar-panel"):
            gr.Markdown("### ⚙️ Engine Configurations")

            model_select = gr.Dropdown(
                choices=["llama-3.3-70b-versatile", "llama-3.1-8b-instant"],
                value="llama-3.3-70b-versatile",
                label="Select AI Model"
            )

            system_input = gr.Textbox(
                value="You are a helpful, brilliant, and concise AI assistant.",
                label="System Prompt / AI Persona",
                lines=3,
                placeholder="Ex: Act as a cynical senior developer..."
            )

            gr.Markdown("---")
            gr.Markdown("### 🧠 Hyperparameters")

            temp_slider = gr.Slider(
                minimum=0.0, maximum=2.0, value=0.7, step=0.1,
                label="Temperature", info="Higher = more creative, Lower = more factual"
            )

            tokens_slider = gr.Slider(
                minimum=128, maximum=4096, value=1024, step=128,
                label="Max Output Tokens"
            )

            gr.Markdown("---")
            gr.Markdown("**Status:** 🟢 Connected via ZeroGPU to Groq")

        # --- RIGHT COLUMN: CHAT INTERFACE ---
        with gr.Column(scale=3):
            # Removed the `type="messages"` keyword argument entirely to prevent the TypeError
            chatbot = gr.Chatbot(elem_classes="chat-window")

            msg_input = gr.Textbox(
                placeholder="Type your message here and press Enter...",
                show_label=False,
                container=False
            )
            
            clear_btn = gr.Button("🗑️ Clear Conversation")

            # Native component wiring
            msg_input.submit(
                fn=chat_stream,
                inputs=[msg_input, chatbot, model_select, system_input, temp_slider, tokens_slider],
                outputs=[chatbot]
            ).then(
                fn=lambda: "", 
                inputs=None, 
                outputs=[msg_input]
            )
            
            # Clear chat utility logic
            clear_btn.click(fn=lambda: [], inputs=None, outputs=[chatbot])

if __name__ == "__main__":
    demo.launch(css=custom_css, theme=gr.themes.Soft())