from openai import OpenAI, AssistantEventHandler
from dotenv import load_dotenv
import os
import gradio as gr
import hashlib

# Load environment variables from .env file
load_dotenv()

# Load env variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ASSISTANT_ID = os.getenv("ASSISTANT_ID")
PASSWORD_HASH = os.getenv("PASSWORD_HASH")  # Stored password hash in .env

client = OpenAI(api_key=OPENAI_API_KEY)

def create_thread():
    return client.beta.threads.create()

def predict(user_message, history, thread):
    # Send the user message to the OpenAI API
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=user_message
    )

    response = []
    # Stream the assistant's response
    with client.beta.threads.runs.stream(
        thread_id=thread.id,
        assistant_id=ASSISTANT_ID
    ) as stream:
        for event in stream:
            if event.event == "thread.message.delta" and event.data.delta.content:
                assistant_message = event.data.delta.content[0].text
                response.append(assistant_message.value)
                yield [{"role": "assistant", "content": ''.join(response)}], thread
    
    # # Append the assistant's response to the history
    history.append({"role": "assistant", "content": ''.join(response)})
    yield history, thread

def verify_password(password):
    """Verify password against the stored hash"""
    if not PASSWORD_HASH:
        return False, "Password not configured in environment variables."
    
    # Hash the input password
    hashed_input = hashlib.sha256(password.encode()).hexdigest()
    print(hashed_input)
    # Compare with stored hash
    if hashed_input == PASSWORD_HASH:
        return True, "Authentication successful"
    else:
        return False, "Incorrect password. Please try again."

# Launch the Gradio chat interface
with gr.Blocks(title="TM Tender Assistant") as demo:
    # Login state
    is_authenticated = gr.State(False)
    
    # Login component
    with gr.Group(visible=True) as login_group:
        gr.Markdown("## TM Tender Assistant - Login")
        password_input = gr.Textbox(type="password", label="Password", placeholder="Enter password")
        login_button = gr.Button("Login")
        login_message = gr.Markdown("")
    
    # Chat interface component
    with gr.Group(visible=False) as chat_group:
        gr.Markdown("## TM Tender Assistant")
        chatbot = gr.Chatbot(type='messages', label="TM Tender Assistant")
        msg = gr.Textbox(placeholder="Type your message here...")
        user_content = gr.State("")
        thread = gr.State(create_thread)
        history = gr.State([])
        submit = gr.Button("Submit")
    
    def login(password):
        success, message = verify_password(password)
        if success:
            return {
                login_group: gr.Group(visible=False),
                chat_group: gr.Group(visible=True),
                login_message: ""
            }
        else:
            print("Fail")
            return {
                login_group: gr.Group(visible=True),
                chat_group: gr.Group(visible=False),
                login_message: ""
            }
    
    login_button.click(
        login,
        inputs=[password_input],
        outputs=[login_group, chat_group, login_message]
    )

    password_input.submit(
        login,
        inputs=[password_input],
        outputs=[login_group, chat_group, login_message]
    )
    
    def user_input(user_message, history):
        user_content = user_message
        return "", history + [{"role": "user", "content": user_message}], user_content

    msg.submit(user_input, [msg, chatbot], [msg, chatbot, user_content], queue=False).then(
        predict, [user_content, chatbot, thread], [chatbot, thread]
    )
    submit.click(user_input, [msg, chatbot], [msg, chatbot, user_content], queue=False).then(
        predict, [user_content, chatbot, thread], [chatbot, thread]
    )

demo.launch(share=False)