import gradio as gr
import random
from helper import load_world, save_world
from together import Together
from helper import get_together_api_key
from guardrail import is_safe
from randomselection import randomly_select_from_json
# Initialize Together client
client = Together(api_key=get_together_api_key())

# Load the world data from the JSON file
world_data = load_world('./YourWorld_L1.json')


# Function to randomly select a world, kingdom, town, and character

# Function to initialize or reinitialize the game state
def initialize_game_state():
    random_state = randomly_select_from_json(world_data)
    world = random_state["world"]
    kingdom = random_state["kingdom"]
    town = random_state["town"]
    character = random_state["character"]

    system_prompt = """You are an AI Game master. Your job is to create a 
    start to an adventure based on the world, kingdom, town, and character 
    a player is playing as. 
    Instructions:
    - You must use only 2-4 sentences.
    - Please use simple and clear language that is easy for children to understand.
    - Write in second person, e.g., "You are Jack."
    - Write in present tense, e.g., "You stand at..."
    - First describe the character and their backstory.
    - Then describe where they start and what they see around them."""

    world_info = f"""
    World: {world['description']}
    Kingdom: {kingdom['description']}
    Town: {town['description']}
    Your Character: {character['description']}
    """

    model_output = client.chat.completions.create(
        model="meta-llama/Llama-3-70b-chat-hf",
        temperature=1.0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": world_info + '\nYour Start:'}
        ],
    )

    start = model_output.choices[0].message.content

    return {
        "world": world["description"],
        "kingdom": kingdom["description"],
        "town": town["description"],
        "character": character["description"],
        "start": start,
    }

    
# Initialize the game state
game_state = initialize_game_state()
game_running = False  # Flag to manage game status

# Function to process user input and actions
def run_action(message, history):
    global game_state, game_running  # Access global variables

    if message.lower() == "restart the game":
        game_state = initialize_game_state()
        game_running = True  # Set the game as running
        return "Game restarted! " + game_state["start"]

    if message.lower() == "start game":
        game_running = True  # Set the game as running
        return game_state["start"]

    if not game_running:
        return "The game has ended. Type 'restart the game' to play again."

    if message.lower() == "exit":
        game_running = False
        return "The game has ended. Type 'restart the game' to play again."

    # The rest of the function remains the same
    system_prompt = """You are an AI Game master. Your job is to write what \
    happens next in a player's adventure game. \
    Instructions: \
    - Write only 1-3 sentences. \
    - Please use simple and clear language that is easy for children to understand. \
    - Always write in second person, e.g., 'You look north and see...' \
    - Write in present tense."""

    world_info = f"""
    World: {game_state['world']}
    Kingdom: {game_state['kingdom']}
    Town: {game_state['town']}
    Your Character: {game_state['character']}"""

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": world_info},
    ]

    for action in history:
        if isinstance(action, tuple) and len(action) == 2:
            messages.append({"role": "assistant", "content": action[0]})
            messages.append({"role": "user", "content": action[1]})

    messages.append({"role": "user", "content": message})

    model_output = client.chat.completions.create(
        model="meta-llama/Llama-3-70b-chat-hf",
        messages=messages,
    )

    return model_output.choices[0].message.content


def main_loop(message, history):

    if not is_safe(message):
        return 'Invalid action.'
    
    result = run_action(message, history)
    safe = is_safe(result)
    if(safe):
        return result # only if safe?
    else:
        return 'Invalid output.'


# Gradio ChatInterface
demo = gr.ChatInterface(
    main_loop,
    chatbot=gr.Chatbot(
        height=450,
        placeholder="Type 'start game' to begin, 'restart the game' to restart, or 'exit' to end the game.",
        type="messages",  # Ensures proper rendering
    ),
    textbox=gr.Textbox(
        placeholder="What do you do next?",
        container=False,
        scale=7,
    ),
    title="AI RPG",
    theme="Monochrome",
    examples=["Look around", "Continue the story"],
    cache_examples=False,
)

# Launch the Gradio app
demo.launch(share=True, server_name="0.0.0.0")