import os
import gradio as gr
from openai import AsyncOpenAI  #Use the Async client
from dotenv import load_dotenv
import re

load_dotenv()

# Global client initialization (Prevents connection overhead)
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    raise RuntimeError("GROQ_API_KEY is not set")

client = AsyncOpenAI(
    api_key=api_key,
    base_url="https://api.groq.com/openai/v1",
    timeout=120.0
)

def sanitize_messages(history):
        sanitized = []
        for m in history:
            # Check if it's the new dictionary format
            if isinstance(m, dict):
                sanitized.append({"role": m["role"], "content": m["content"]})
            # Check if it's the old list/tuple format [user, bot]
            elif isinstance(m, (list, tuple)):
                if m[0]: sanitized.append({"role": "user", "content": m[0]})
                if m[1]: sanitized.append({"role": "assistant", "content": m[1]})
        return sanitized

def strip_think(text: str) -> str:
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

async def respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    
    messages = [{"role": "system", "content": system_message}]
    messages.extend(sanitize_messages(history))
    messages.append({"role": "user", "content": message})

    try:
        completion = await client.chat.completions.create(
            model="qwen/qwen3-32b",
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )

        text = completion.choices[0].message.content
        return strip_think(text)

    except Exception as e:
        return f"Error: {str(e)}"


chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

if __name__ == "__main__":
    chatbot.queue(default_concurrency_limit=5).launch()