import gradio as gr
from huggingface_hub import InferenceClient
import re
import traceback
import os

# Initialize the Inference Client with your model
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=os.getenv("HUGGINGFACE_API_TOKEN"))

def respond(
    message,
    history,
    max_tokens,
    temperature,
    top_p,
    current_salary
):
    """
    Respond function to handle the conversation and update salary based on AI's assessment.
    """
    # Define the system message for the conversation (hidden from the user)
    system_message = (
        "You are a hiring manager negotiating a job offer with a candidate. "
        "Your initial salary offer is $60,000. "
        "Engage in a negotiation with the candidate, adjusting your offer based on their arguments."
    )
    
    # Initialize the messages with the system prompt
    messages = [{"role": "system", "content": system_message}]
    
    # Append the conversation history to messages
    for user_msg, ai_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if ai_msg:
            messages.append({"role": "assistant", "content": ai_msg})
    
    # Append the latest user message
    messages.append({"role": "user", "content": message})
    
    # Generate the AI's response to the user's message
    try:
        response = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=None
        ).get("choices")[0].get("message").get("content")
    except Exception as e:
        response = f"An error occurred while communicating with the AI: {e}"
        traceback.print_exc()
    
    # Append the AI's response to the history
    history.append((message, response))
    
    # Now, send the conversation to the AI to get the updated salary
    # Prepare the salary assessment prompt
    salary_assessment_prompt = (
        "As the hiring manager, based on the conversation so far, "
        "what salary do you now offer to the candidate? "
        "Please provide only the salary amount as a number, without any additional text."
    )
    
    # Prepare the messages for salary assessment
    assessment_messages = [{"role": "system", "content": salary_assessment_prompt}]
    
    # Include the conversation history
    for user_msg, ai_msg in history:
        if user_msg:
            assessment_messages.append({"role": "user", "content": user_msg})
        if ai_msg:
            assessment_messages.append({"role": "assistant", "content": ai_msg})
    
    # Generate the AI's salary assessment
    try:
        salary_response = client.chat_completion(
            assessment_messages,
            max_tokens=10,
            temperature=temperature,
            top_p=top_p,
            stop=None
        ).get("choices")[0].get("message").get("content")
    except Exception as e:
        salary_response = f"An error occurred while assessing salary: {e}"
        traceback.print_exc()
        salary_response = str(current_salary)
    
    # Parse the salary from the AI's response
    try:
        # Remove any non-digit characters
        salary_str = re.sub(r'[^\d]', '', salary_response)
        if salary_str:
            new_salary = int(salary_str)
            if new_salary != current_salary:
                # Update current_salary
                current_salary = new_salary
        else:
            # If parsing fails, keep the current salary
            pass
    except Exception as e:
        # If parsing fails, keep the current salary
        pass
    
    return history, "", current_salary  # Return history, clear input, and update salary internally

def reset_game():
    """
    Function to reset the game to initial state.
    """
    return [], 60000  # Reset history and salary to $60,000

# Define the Gradio Blocks layout
with gr.Blocks() as demo:
    gr.Markdown("# 💼 Salary Negotiation Game")
    gr.Markdown(
        """
        **Objective:** Negotiate your salary starting from $60,000.

        **Instructions:**
        - Use the chat to negotiate your salary with the hiring manager.
        - Provide compelling reasons for a higher salary.
        - The hiring manager will consider your arguments and may adjust the offer.
        """
    )

    # Chat history to keep track of the conversation
    chat_history = gr.State([])

    # Current salary (hidden from the user)
    current_salary_state = gr.State(60000)

    # Chat Interface
    chatbot = gr.Chatbot()

    # User input
    user_input = gr.Textbox(
        label="Your Message",
        placeholder="Enter your negotiation message here...",
        lines=2
    )
    send_button = gr.Button("Send")

    def handle_message(message, history, max_tokens, temperature, top_p, current_salary):
        """
        Handles user messages and updates the conversation history and salary.
        """
        history, _, current_salary = respond(message, history, max_tokens, temperature, top_p, current_salary)
        return history, "", current_salary

    send_button.click(
        handle_message,
        inputs=[
            user_input,
            chat_history,
            gr.Number(value=512, label="Max New Tokens"),
            gr.Number(value=0.7, label="Temperature"),
            gr.Number(value=0.95, label="Top-p"),
            current_salary_state  # Pass the current salary
        ],
        outputs=[
            chatbot,
            user_input,          # Clear the input textbox
            current_salary_state  # Update the salary internally
        ]
    )

    # Reset button to restart the game
    reset_btn = gr.Button("Reset Game")
    reset_btn.click(fn=reset_game, inputs=None, outputs=[chat_history, current_salary_state])

    gr.Markdown(
        """
        ---
        *Developed with ❤️ using Gradio and Hugging Face.*
        """
    )

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