import gradio as gr
from groq import Groq
import os
import time
import json

# =========================
# API KEY
# =========================
api_key = os.getenv("GROQ_API_KEY")

if not api_key:
    raise ValueError("GROQ_API_KEY not found")

client = Groq(api_key=api_key.strip())

# =========================
# MEMORY FILE
# =========================
MEMORY_FILE = "memory.json"

def load_memory():
    if os.path.exists(MEMORY_FILE):
        with open(MEMORY_FILE, "r") as f:
            return json.load(f)
    return {"facts": {}}

def save_memory(data):
    with open(MEMORY_FILE, "w") as f:
        json.dump(data, f)

memory = load_memory()

# =========================
# CHAT FUNCTION
# =========================
def chat(message, history):

    global memory

    # -------------------------
    # MEMORY DETECTION
    # -------------------------
    if "my name is" in message.lower():
        name = message.lower().split("my name is")[-1].strip()
        memory["facts"]["name"] = name
        save_memory(memory)

    # -------------------------
    # SYSTEM PROMPT WITH MEMORY
    # -------------------------
    system_prompt = "You are a helpful AI assistant."

    if "name" in memory["facts"]:
        system_prompt += f" User name is {memory['facts']['name']}."

    messages = [{"role": "system", "content": system_prompt}]

    # -------------------------
    # SAFE HISTORY HANDLING (FIXED ERROR)
    # -------------------------
    for item in history:
        if isinstance(item, (list, tuple)) and len(item) == 2:
            user_msg, bot_msg = item
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": bot_msg})

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

    # -------------------------
    # STREAM RESPONSE
    # -------------------------
    stream = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        stream=True,
        temperature=0.7
    )

    response_text = ""

    for chunk in stream:
        if chunk.choices[0].delta.content:
            token = chunk.choices[0].delta.content
            response_text += token

            time.sleep(0.03)

            yield response_text + "▌"

    yield response_text

# =========================
# UI
# =========================
demo = gr.ChatInterface(
    fn=chat,
    title="Streaming AI Chatbot using Groq",
    description="AI chatbot with memory + streaming",
)

demo.launch()