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

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

# Load the model and tokenizer from Hugging Face
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-mini-instruct",
    device_map="cpu",
    # device_map="cuda",
    torch_dtype="auto",
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-mini-instruct")

# Create a text-generation pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

# Define the pipeline arguments
generation_args = {
    "max_new_tokens": 500,
    "return_full_text": False,
    "temperature": 0.0,
    "do_sample": False,
}

chat_session = {}


# Function to generate responses based on the entire chat history
def generate_response(chat_history):
    messages = [{"role": "system", "content": "You are a helpful AI assistant."}]
    
    # Append the chat history (user and assistant messages)
    for user_message, assistant_message in chat_history:
        messages.append({"role": "user", "content": user_message})
        messages.append({"role": "assistant", "content": assistant_message})
    
    # Generate response for the latest user message
    user_input = chat_history[-1][0]  # The latest user message
    messages.append({"role": "user", "content": user_input})
    response = pipe(messages, **generation_args)
    
    # Append the assistant response to the chat history
    assistant_response = response[0]["generated_text"]
    return assistant_response


# Function to update chat
def chat(user_message,history,session):

    if session == "":
        return history, "Error: Session ID cannot be empty. Please start a new chat."
        

    history = history or []  # Initialize history if empty
        
    # Generate assistant response based on the history
    assistant_message = generate_response(history + [(user_message, "")])
    
    # Append user and assistant messages to history
    history.append([user_message, assistant_message])
    chat_session[session] = history

    print("USER : ",user_message)
    print("ASSISTANT : ",assistant_message)

    return history, ""


def get_session_list():
    return list(chat_session.keys())

# Function to create new chat and return updated session list
def new_chat():
    session = f'session:{len(chat_session) + 1}'
    chat_session[session] = []  # Initialize empty chat history for the new session
    return [], "", session, get_session_list()  # Return the new session and update session list

# Function to fetch old chat session history
def old_chat(sessions):
    return chat_session.get(sessions, [])

# Function to reset chat history
def reset_button():
    global chat_session  # Access the global chat_session
    chat_session = {}  # Reset the global chat_session
    return [], "", "",[],""  # Reset chat history, session, and input field



with gr.Blocks(css=".small-btn {width: 100px !important;} .large-textbox {width: 100% !important;}") as demo:
    gr.Markdown("# 🤖 AI Assistant")

    # Add instructions to the UI
    gr.Markdown("""
    ## Steps to Use the AI Assistant:

    1. **Start a New Chat**: Click the **'Start New Chat'** button to create a new session.
    2. **Send a Message**: Type your message in the input box and either press **Enter** or click **'Send'** to get a response from the AI.
    3. **View Available Sessions**: Click **'Get Available Session'** to list all your chat sessions.
    4. **Load an Old Session**: You can enter the previous session ID in the **'Load Session'** box and click **'Load Session'**.
    5. **Reset All Chats**: Click the **'Reset All'** button to clear all chat sessions and start fresh.
    
    **Session ID**: Every new chat has a unique session ID, which you can use to return to previous conversations.
    """)

    with gr.Column():
        new_chat_button = gr.Button("Start New Chat")
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(elem_id="chatbot")
            with gr.Row():
                with gr.Column(scale=5):
                    user_input = gr.Textbox(
                        show_label=False,
                        placeholder="Type your message here...",
                        container=False,
                        elem_classes="large-textbox"
                    )
                with gr.Column(scale=1):
                    send_button = gr.Button("Send", variant="primary", elem_classes="small-btn")
            
        with gr.Column(scale=1):
            session = gr.Textbox(label="Current Session", interactive=False)
            session_list = gr.Dropdown(label="Available Sessions", choices=get_session_list(), allow_custom_value=True)
            load_session = gr.Textbox(label="Load Session", interactive=True)
            with gr.Row():
                get_old_session_button = gr.Button("Load Session")
                avail_session = gr.Button("Get Available Session")
            reset_button_ = gr.Button("Reset All", variant="secondary")

    # Button click actions
    user_input.submit(chat, [user_input, chatbot, session], [chatbot, user_input])
    send_button.click(chat, [user_input, chatbot, session], [chatbot, user_input])  # Send button
    new_chat_button.click(new_chat, [], [chatbot, user_input, session, session_list])  # Also update the session list
    get_old_session_button.click(old_chat, [load_session], [chatbot])
    reset_button_.click(reset_button, [], [chatbot, session, user_input, session_list, load_session])
    avail_session.click(get_session_list, [], [session_list])

# Launch the Gradio app
demo.launch()