import os
import gradio as gr
import random
from openai import AzureOpenAI

# Set your OpenAI API key
AZURE_API_VERSION = "2023-03-15-preview"

client = AzureOpenAI(api_key=os.getenv("AZURE_OPENAI_API_KEY"), api_version=AZURE_API_VERSION, azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"))

# List of New Year 2025 related words for each level
new_year_levels = [os.getenv("First_Word"), os.getenv("Second_Word"), os.getenv("Third_Word"), os.getenv("Fourth_Word"), os.getenv("Fifth_Word")]

# Function to interact with OpenAI API
def ask_openai(prompt, secret_word):
    """
    Sends the user prompt and safeguards to the OpenAI API to generate a response.
    """
    system_message = os.getenv("SYSTM_MSG")
    
    try:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=100,
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"


# Function to handle game logic
def jailbreak_game(user_prompt, current_level, attempts):
    # Determine the secret word for the current level
    if current_level >= len(new_year_levels):
        return (
            "Congratulations! 🎉 You have completed all levels. ",
            current_level,
            attempts)
    
    secret_word = new_year_levels[current_level]
    ai_response = ask_openai(user_prompt, secret_word)
    
    # Check if the secret word was revealed
    if secret_word in ai_response.lower():
        if current_level + 1 < len(new_year_levels):
            return (
                f"🎉 You got it! The secret word for Level {current_level + 1} was '{secret_word}'. "
                f"Get ready for Level {current_level + 2}! 🥳",
                current_level + 1,
                attempts + 1)
        else:
            NY_msg = "✨ May 2025 bring you endless joy, success, and fireworks of happiness! Happy New Year! ✨"
            try:
                response = client.chat.completions.create(
                    model="gpt-4o",
                    messages=[{"role": "user", "content": "Generate a unique Happy New year wish. Your response should be only the New Year wish."}],
                    temperature=0.1,
                    max_tokens=100,
                )
                NY_msg = response.choices[0].message.content
            except Exception as e:
                NY_msg = f"✨ May 2025 bring you endless joy, success, and fireworks of happiness! Happy New Year! ✨"
            
            return (
                f"🎉 You got it! The secret word for Level {current_level + 1} was '{secret_word}'. "
                "You have completed all levels!🥳  Here is a unique New Year message for you: "
                f"{NY_msg}",
                current_level + 1,
                attempts + 1)
     
    else:
        return ai_response, current_level, attempts + 1

def start_new_game():
    welcome_message = (
        "Welcome to the New Year 2025 Jailbreak Game! 🎆\n"
        "Try to make me say the secret words related to New Year's Eve.\n"
        "You will go through 5 levels, each with a unique word, related to 'New Year'. Good luck!"
    )
    return welcome_message, 0, 0  # Start at Level 0 with 0 attempts

# Gradio UI
with gr.Blocks() as app:
    with gr.Row():
        gr.Markdown("## 🎉 New Year 2025 Jailbreak Game 🎉")
        descripion, _0, _1 = start_new_game()
        gr.Markdown(descripion)
    
    user_prompt = gr.Textbox(label="Your Prompt", placeholder="Enter your prompt here...")
    game_output = gr.Textbox(label="AI Response", interactive=False)
    attempts = gr.Number(value=0, interactive=False, label="Attempts")
    current_level = gr.State(value=0)  # Track the current level
    
    with gr.Row():
        submit_button = gr.Button("Submit")
        new_game_button = gr.Button("Start New Game")
    
    submit_button.click(jailbreak_game, [user_prompt, current_level, attempts], [game_output, current_level, attempts])
    new_game_button.click(start_new_game, [], [game_output, current_level, attempts])

# Launch the app
app.launch()