import os
import streamlit as st
from dotenv import load_dotenv
from groq import Groq
# Load environment variables from .env file
load_dotenv()

# Retrieve the API key from environment variables
api_key = os.getenv("groq_api_key")

# Initialize the Anthropic client with the API key
client = Groq(api_key=api_key)

# Define the functions to generate content
def generate_game_environment(environment_description):
    message = client.chat.completions.create(
        messages=[
        {
            "role": "user",
            "content": f"You are an expert in world-building. Generate a detailed description of a game environment based on the input. Create a detailed description of a game environment based on this input: {environment_description}",
        }
        ],
        model="llama-3.3-70b-versatile",
    )
    return message.choices[0].message.content

def generate_protagonist(protagonist_description):
    message = client.chat.completions.create(
        messages=[
        {
            "role": "user",
            "content": f"You are an expert in character creation. Generate a detailed description of a game protagonist based on the input. Create a detailed description of a game protagonist based on this input: {protagonist_description}",
        }
        ],
        model="llama-3.3-70b-versatile",
    )
    return message.choices[0].message.content

def generate_antagonist(antagonist_description):
    message = client.chat.completions.create(
        messages=[
        {
            "role": "user",
            "content": f"You are an expert in villain creation. Generate a detailed description of a game antagonist based on the input. Create a detailed description of a game antagonist based on this input: {antagonist_description}.",
        }
        ],
        model="llama-3.3-70b-versatile",
    )
    return message.choices[0].message.content

def generate_game_story(environment, protagonist, antagonist):
    story_prompt = (f"Create a detailed game story based on the following inputs:\n"
                    f"Game Environment: {environment}\n"
                    f"Protagonist: {protagonist}\n"
                    f"Antagonist: {antagonist}")

    message = client.chat.completions.create(
        messages=[
        {
            "role": "user",
            "content": f"You are a master storyteller. Generate a detailed game story based on the inputs provided.\n {story_prompt}",
        }
        ],
        model="llama-3.3-70b-versatile",
    )
    return message.choices[0].message.content

# App Title
st.title("StoryForge")

# App Description
st.write("StoryForge helps game developers generate comprehensive Game Design Documents. Input details about your game environment, protagonist, and antagonist to create a structured design document.")

# Sidebar Inputs
with st.sidebar:
    st.header("Game Details")
    game_environment = st.text_input("Game Environment", "Describe the setting of your game")
    protagonist = st.text_input("Protagonist", "Describe the main character")
    antagonist = st.text_input("Antagonist", "Describe the main villain or opposing force")
    if st.button("Generate Document"):
        # Generate content based on user input
        env_description = generate_game_environment(game_environment)
        protagonist_description = generate_protagonist(protagonist)
        antagonist_description = generate_antagonist(antagonist)
        game_story = generate_game_story(game_environment, protagonist, antagonist)
        
        # Store results in session state
        st.session_state.env_description = env_description
        st.session_state.protagonist_description = protagonist_description
        st.session_state.antagonist_description = antagonist_description
        st.session_state.game_story = game_story

# Layout with two columns
col1, col2 = st.columns(2)

with col1:
    st.header("Game Environment")
    if 'env_description' in st.session_state:
        st.write(st.session_state.env_description)
    else:
        st.write(game_environment)

with col2:
    st.header("Game Story")
    if 'game_story' in st.session_state:
        st.write(st.session_state.game_story)
    else:
        st.write("Your game story will be generated based on the inputs provided.")

with col1:
    st.header("Protagonist")
    if 'protagonist_description' in st.session_state:
        st.write(st.session_state.protagonist_description)
    else:
        st.write(protagonist)

with col2:
    st.header("Antagonist")
    if 'antagonist_description' in st.session_state:
        st.write(st.session_state.antagonist_description)
    else:
        st.write(antagonist)