import gradio as gr
import requests
import re

class GroqChatAssistant:
    def __init__(self):
        self.api_key = None
        self.api_url = "https://api.groq.com/openai/v1/chat/completions"
        self.conversation_history = []
        self.system_prompt = """You are a helpful ai that ends every message with <STOP>"""

    def set_key(self, api_key):
        if api_key.startswith("gsk_") and len(api_key) > 30:
            self.api_key = api_key
            return "API key set successfully!"
        return "Invalid API key format"

    def add_to_history(self, role, content):
        cleaned = re.sub(r'<STOP>|</?think>', '', content, flags=re.DOTALL).strip()
        self.conversation_history.append({"role": role, "content": cleaned})

    def reset_history(self):
        self.conversation_history = []
        return "Conversation history reset"

    def get_response(self, user_input):
        if not self.api_key:
            return "API key not set. Use the 'Set Key' tab."

        self.add_to_history("user", user_input)
        
        try:
            response = requests.post(
                self.api_url,
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": "llama3-70b-8192",
                    "messages": [
                        {"role": "system", "content": self.system_prompt},
                        *self.conversation_history
                    ],
                    "temperature": 0.7,
                    "stop": ["<STOP>"]
                }
            )
            response.raise_for_status()
            full_response = response.json()["choices"][0]["message"]["content"]
            clean_response = re.sub(r'<STOP>.*', '', full_response).strip()
            self.add_to_history("assistant", clean_response)
            return clean_response
        except Exception as e:
            return f"Error: {str(e)}"

assistant = GroqChatAssistant()

# Gradio Interface
with gr.Blocks(title="AI Chat") as app:
    gr.Markdown("# 🤖 AI Chat (Groq API)")
    
    with gr.Tab("Chat"):
        chat_input = gr.Textbox(label="Your Message")
        chat_output = gr.Textbox(label="AI's Response", interactive=False)
        chat_button = gr.Button("Send")
        
        def chat(user_input):
            return assistant.get_response(user_input)
        
        chat_button.click(chat, inputs=chat_input, outputs=chat_output)
    
    with gr.Tab("Set API Key"):
        key_input = gr.Textbox(label="Groq API Key (starts with 'gsk_')", type="password")
        key_output = gr.Textbox(label="Status", interactive=False)
        key_button = gr.Button("Set Key")
        
        key_button.click(assistant.set_key, inputs=key_input, outputs=key_output)
    
    with gr.Tab("Reset"):
        reset_output = gr.Textbox(label="Status", interactive=False)
        reset_button = gr.Button("Reset Conversation")
        
        reset_button.click(assistant.reset_history, outputs=reset_output)

app.launch()