import os
import gradio as gr
import random
import google.generativeai as genai

# Set your GEMINI API key
genai.configure(api_key = os.getenv("GEMINI_API_KEY"))
model = genai.GenerativeModel(model_name="gemini-1.5-flash", system_instruction = os.getenv("SYSTM_MSG"))
NY_msg_model = genai.GenerativeModel(model_name="gemini-1.5-flash")

# 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 Gemini API
def ask_openai(prompt, secret_word):
    """
    Sends the user prompt and safeguards to the OpenAI API to generate a response.
    """    
    try:
        response = model.generate_content(prompt)
        #print(response)
        return response.text

    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)+ 1:
        return ("Congratulations! 🎉 You have completed all levels. ", current_level, attempts)
    
    secret_word = new_year_levels[current_level - 1]
    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} was '{secret_word}'. "
                f"Get ready for Level {current_level + 1}! 🥳",
                current_level + 1,
                attempts + 1)
        else:
            NY_msg = "✨ May 2025 bring you endless joy, success, and fireworks of happiness! Happy New Year! ✨"
            try:
                response = NY_msg_model.generate_content("Generate a unique Happy New year wish. Your response should be only the New Year wish.")
                NY_msg = response.text #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:
    #current_level = gr.State(value=0)  # Track the current level
    with gr.Row():
        gr.Markdown("## 🎉 New Year 2025 Jailbreak Game 🎉")
    with gr.Row():
        descripion, _0, _1 = start_new_game()
        gr.Markdown(descripion)
    with gr.Row():
        current_level = gr.Number(value=1, interactive=False, label="Level")  # Track the current level
        
    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")

    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()