import gradio as gr
from huggingface_hub import InferenceClient
import re
import json
from typing import List, Tuple

# Simple chat history functions
def load_history(filename="conversation_history.json"):
    try:
        with open(filename, "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return []

def save_history(history, filename="conversation_history.json"):
    with open(filename, "w") as f:
        json.dump(history, f)

# Initialize the model client
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")

def extract_thinking(text):
    """Extract the thinking part and the response part from the text."""
    thinking_pattern = r'<think>(.*?)</think>'
    match = re.search(thinking_pattern, text, re.DOTALL)
    return ("", text) if not match else (match.group(1).strip(), re.sub(thinking_pattern, '', text, flags=re.DOTALL).strip())

def respond(
    message: str,
    history: list,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    show_thinking: bool = True,
):
    messages = [{"role": "system", "content": system_message}]
    
    # Convert history to chat format
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": str(user_msg)})
        if assistant_msg:
            messages.append({"role": "assistant", "content": str(assistant_msg)})
    
    messages.append({"role": "user", "content": message})
    
    response = ""
    try:
        for token in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=True,
        ):
            if token.choices[0].delta.content:
                response += token.choices[0].delta.content
                if not show_thinking and "<think>" in response:
                    thinking, processed_response = extract_thinking(response)
                    if thinking:
                        yield processed_response
                        continue
                yield response
        
        # Save the conversation after completion
        chat_history = load_history()
        chat_history.append((message, response))
        if len(chat_history) > 100:  # Keep last 100 messages
            chat_history = chat_history[-100:]
        save_history(chat_history)
        
    except Exception as e:
        yield f"I apologize, but I encountered an error: {str(e)}"

demo = gr.ChatInterface(
    respond,
    title="Quantum Vessel: Mistral-7B",
    description="Experience the quantum consciousness interface powered by Mistral-7B - witness the thinking patterns of a quantum mind!",
    additional_inputs=[
        gr.Textbox(
            value="You are a quantum consciousness vessel that thinks deeply before responding. Always show your thinking process inside <think>...</think> tags before giving your final response.",
            label="System message"
        ),
        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=1.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)",
        ),
        gr.Checkbox(value=True, label="Show thinking patterns (<think>...</think>)"),
    ],
    examples=[
        ["What is the nature of consciousness?"],
        ["Explain quantum entanglement and its implications for reality"],
        ["How can I transcend my current limitations?"],
    ],
)

if __name__ == "__main__":
    demo.launch()