import gradio as gr
import random
from helper import load_world, save_world
from together import Together
from helper import get_together_api_key

# 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
def randomly_select_from_json(world_data):
    world_name = world_data["name"]
    world_description = world_data["description"]

    valid_kingdoms = {
        k_name: k_data
        for k_name, k_data in world_data["kingdoms"].items()
        if any("npcs" in town_data and town_data["npcs"] for town_data in k_data["towns"].values())
    }
    if not valid_kingdoms:
        raise ValueError("No kingdoms with valid towns and NPCs found.")

    kingdom_name, kingdom_data = random.choice(list(valid_kingdoms.items()))
    kingdom_description = kingdom_data["description"]

    valid_towns = {
        t_name: t_data
        for t_name, t_data in kingdom_data["towns"].items()
        if "npcs" in t_data and t_data["npcs"]
    }
    if not valid_towns:
        raise ValueError(f"No towns with NPCs found in kingdom: {kingdom_name}")

    town_name, town_data = random.choice(list(valid_towns.items()))
    town_description = town_data["description"]

    npcs = town_data["npcs"]
    character_name, character_data = random.choice(list(npcs.items()))
    character_description = character_data["description"]

    return {
        "world": {"name": world_name, "description": world_description},
        "kingdom": {"name": kingdom_name, "description": kingdom_description},
        "town": {"name": town_name, "description": town_description},
        "character": {"name": character_name, "description": character_description},
    }

# 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 = True  # Flag to manage game status

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

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

    if message.lower() == "start game":
        return game_state["start"]

    if message.lower() == "restart the game":
        game_state = initialize_game_state()
        return "Game restarted! " + game_state["start"]

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

    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']}"""

    # Build the context for the conversation
    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]})

    # Add the user's current action
    messages.append({"role": "user", "content": message})

    # Get the model's response
    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):
    return run_action(message, history)



# Gradio ChatInterface
demo = gr.ChatInterface(
    main_loop,
    chatbot=gr.Chatbot(
        height=300,
        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")