import gradio as gr
import uuid
import json
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the DeepSeek-R1 model
model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-R1",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1")

def format_messages_for_model(messages):
    """
    Format the conversation history into a single string compatible with the DeepSeek-R1 model.
    """
    formatted_messages = ""
    for msg in messages:
        if msg["role"] == "user":
            formatted_messages += f"User: {msg['content']}\n"
        elif msg["role"] == "assistant":
            formatted_messages += f"DeepSeek-R1: {msg['content']}\n"
    return formatted_messages

def generate_response(user_input, messages):
    """
    Generate a response using DeepSeek-R1 and return it as a string.
    """
    formatted_history = format_messages_for_model(messages)
    full_input = f"{formatted_history}User: {user_input}\nDeepSeek-R1: "

    inputs = tokenizer(full_input, return_tensors="np")
    outputs = model.generate(
        **inputs,
        max_length=2048,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        return_dict_in_generate=True,
    )
    
    response = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
    return response.split("DeepSeek-R1: ")[-1].strip()

def chatbot(
    message,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """
    Main function for generating responses and saving chat history.
    """
    # Generate a unique session ID for this chat session
    session_id = str(uuid.uuid4())

    # Initialize the messages with the system message
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})

    # Add previous chat history
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})

    # Add the new user message
    messages.append({"role": "user", "content": message})

    # Generate response
    response = generate_response(message, messages)

    # Add the bot's response to messages
    messages.append({"role": "assistant", "content": response})

    # Update chat history in the interface
    new_history = history + [(message, response)]

    # Save chat history to Hugging Face storage
    storage = gr.HuggingFace()
    storage.set(key=session_id, value={"messages": messages})

    return response, new_history

def create_chat_interface():
    """
    Create the Gradio chat interface.
    """
    with gr.Blocks() as app:
        session_id = str(uuid.uuid4())  # Generate session ID at startup

        with gr.Row():
            chatbot_interface = gr.Chatbot(label="DeepSeek-R1 Chatbot")
            input_box = gr.Textbox(
                label="Type your message here",
                placeholder="Type your message here...",
                show_label=True,
            )

        # Button to submit the message
        submit_button = gr.Button("Send")

        # Store chat history in the interface
        history = gr.State([])

        # Define the system message
        system_message = gr.Textbox(
            value="You are DeepSeek-R1, an AI assistant.",
            label="System Message",
            show_label=True,
            visible=True,
        )

        # Define inference parameters
        with gr.Accordion("Parameters", open=True):
            max_tokens = gr.Slider(
                minimum=1,
                maximum=2048,
                value=512,
                step=1,
                label="Max Tokens",
                info="Maximum length of the generated response.",
            )
            temperature = gr.Slider(
                minimum=0.0,
                maximum=2.0,
                value=1.0,
                step=0.1,
                label="Temperature",
                info="Higher values produce more diverse outputs.",
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.9,
                step=0.05,
                label="Top-p",
                info="Controls diversity through nucleus sampling.",
            )

        # Define the submit function
        submit_button.click(
            fn=chatbot,
            inputs=[
                input_box,
                history,
                system_message,
                max_tokens,
                temperature,
                top_p,
            ],
            outputs=[gr.Textbox(), chatbot_interface],
            api_name="chatbot",
        )

    return app

if __name__ == "__main__":
    app = create_chat_interface()
    app.launch()