import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Set seed for reproducibility
torch.random.manual_seed(0)

# Load the model and tokenizer (using CPU)
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-mini-instruct",
    device_map="cpu",  # Use CPU
    torch_dtype="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-mini-instruct")

# Define the pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=-1  # CPU
)

# System message (invisible to the user)
SYSTEM_MESSAGE = {"role": "system", "content": "You are a helpful AI assistant."}

# Function to process the user input and generate output
def chatbot_response(conversation_history):
    # Build message sequence
    messages = [SYSTEM_MESSAGE] + [
        {"role": "user", "content": message["user_input"]} for message in conversation_history
    ]
    # Pass messages to the model
    generation_args = {
        "max_new_tokens": 500,
        "return_full_text": False,
        "temperature": 0.0,
        "do_sample": False,
    }
    output = pipe(messages, **generation_args)
    assistant_reply = output[0]["generated_text"]
    # Append assistant's response to history
    conversation_history[-1]["assistant_reply"] = assistant_reply
    return conversation_history

# Define Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# AI Chatbot with System Message")
    
    with gr.Row():
        with gr.Column():
            chatbox = gr.Chatbot(label="Conversation")
            input_box = gr.Textbox(label="Your Message", placeholder="Type your message here...")
            submit_btn = gr.Button("Submit")
    
    conversation_state = gr.State([])  # Maintain conversation history

    # Function to update the conversation
    def update_conversation(user_input, history):
        if user_input.strip():
            history.append({"user_input": user_input})
            updated_history = chatbot_response(history)
            # Format conversation history for the Chatbot component
            formatted_conversation = [
                (msg["user_input"], msg.get("assistant_reply", ""))
                for msg in updated_history
            ]
            return formatted_conversation, updated_history, ""
        return [], history, ""

    submit_btn.click(
        update_conversation,
        inputs=[input_box, conversation_state],
        outputs=[chatbox, conversation_state, input_box],
    )

# Launch the interface
demo.launch()