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

client = Together(api_key=get_together_api_key())

 world = load_world('./YourWorld_L1.json')
 kingdom = world['kingdoms']['Varanasi Kingdom']
 town = kingdom['towns']["Sarnath"]
 character = town['npcs']['Rohan Vyas']

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 only use 2-4 sentences \
Write in second person. For example: "You are Jack" \
Write in present tense. For example "You stand at..." \
First describe the character and their backstory. \
Then describes where they start and what they see around them."""

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

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
print(start)
world['start'] = start
#save_world(world, '../shared_data/Kyropeia.json')  # preserve video version
save_world(world, '.YourWorld_L1.json')

game_state = {
    "world": world['description'],
    "kingdom": kingdom['description'],
    "town": town['description'],
    "character": character['description'],
    "start": start,
}

def run_action12(message, history, game_state):
    
    if(message == 'start game'):
        return game_state['start']

    system_prompt = """You are an AI Game master. Your job is to write what \
happens next in a player's adventure game.\
Instructions: \
You must on only write 1-3 sentences in response. \
Always write in second person present tense. \
Ex. (You look north and see...)"""
    
    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:
        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
    )
    
    result = model_output.choices[0].message.content
    return result

def run_action(message, history, game_state):
    if message.lower() == 'start game':
        return game_state['start']

    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. \
- 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 conversation context
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": world_info}
    ]

    for action in history:
        # Ensure action is in the correct format before appending
        if isinstance(action, tuple) and len(action) == 2:
            messages.append({"role": "assistant", "content": action[0]})
            messages.append({"role": "user", "content": action[1]})
        else:
            # If action is not in the correct format, log a warning or handle the case
            print(f"Warning: Invalid action format: {action}")

    # Add the user's current message
    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, game_state)

demo = None  # Added to allow restart

def start_game(main_loop, share=False):
    global demo
    # If demo is already running, close it first
    if demo is not None:
        demo.close()

    demo = gr.ChatInterface(
        main_loop,
        chatbot=gr.Chatbot(
            height=250, 
            placeholder="Type 'start game' to begin", 
            type='messages'  # Updated parameter for Gradio's chatbot messages
        ),
        textbox=gr.Textbox(
            placeholder="What do you do next?", 
            container=False, 
            scale=7
        ),
        title="AI RPG",
        theme="soft",
        examples=["Look around", "Continue the story"],
        cache_examples=False,
    )
    demo.launch(share=share, server_name="0.0.0.0")


start_game(main_loop, True)