import os
from dotenv import load_dotenv
import gradio as gr
import anthropic

# Load environment variables from .env file
load_dotenv()

# Get the API key from environment variables
def get_client():
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        raise ValueError("API key not found. Please set the ANTHROPIC_API_KEY environment variable.")
    return anthropic.Anthropic(api_key=api_key)

def generate_game_environment(description_prompt):
    client = get_client()
    response = client.messages.create(
        model="claude-3-5-sonnet-20240620",
        max_tokens=500,
        temperature=0.7,
        system="You are a creative writer specialized in generating immersive game environments.",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": description_prompt
                    }
                ]
            }
        ]
    )
    return response['choices'][0]['text'].strip()

def generate_protagonist(description_prompt):
    client = get_client()
    response = client.messages.create(
        model="claude-3-5-sonnet-20240620",
        max_tokens=500,
        temperature=0.7,
        system="You are a character designer specializing in creating detailed protagonists for games.",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": description_prompt
                    }
                ]
            }
        ]
    )
    return response['choices'][0]['text'].strip()

def generate_antagonist(description_prompt):
    client = get_client()
    response = client.messages.create(
        model="claude-3-5-sonnet-20240620",
        max_tokens=500,
        temperature=0.7,
        system="You are a character designer specializing in creating compelling antagonists for games.",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": description_prompt
                    }
                ]
            }
        ]
    )
    return response['choices'][0]['text'].strip()

def generate_game_story(environment, protagonist, antagonist):
    description_prompt = (
        f"Create a compelling game story based on the following details:\n"
        f"Game Environment: {environment}\n"
        f"Protagonist: {protagonist}\n"
        f"Antagonist: {antagonist}\n"
        "Please provide a narrative that ties these elements together in an engaging way."
    )
    client = get_client()
    response = client.messages.create(
        model="claude-3-5-sonnet-20240620",
        max_tokens=1000,
        temperature=0.7,
        system="You are a master storyteller. Generate an engaging game story from the given details.",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": description_prompt
                    }
                ]
            }
        ]
    )
    return response['choices'][0]['text'].strip()

def generate_game_design(environment, protagonist, antagonist):
    # Generate each part of the game design
    env_description = generate_game_environment(environment)
    prot_description = generate_protagonist(protagonist)
    antag_description = generate_antagonist(antagonist)
    game_story = generate_game_story(env_description, prot_description, antag_description)
    
    return env_description, prot_description, antag_description, game_story

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# StoryForge\nA tool to help you generate engaging game design ideas.")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("## Game Development")
                environment = gr.Textbox(label="Game Environment", placeholder="e.g., a post-apocalyptic city")
                protagonist = gr.Textbox(label="Protagonist", placeholder="e.g., a young warrior with a mysterious past")
                antagonist = gr.Textbox(label="Antagonist", placeholder="e.g., a dark sorcerer seeking to conquer the world")
                submit_btn = gr.Button("Generate Game Design")
            
            with gr.Column():
                gr.Markdown("## Protagonist")
                gr.Markdown("## Antagonist")
                gr.Markdown("## Game Story")
                env_output = gr.Textbox(label="Game Environment", interactive=False)
                prot_output = gr.Textbox(label="Protagonist", interactive=False)
                antag_output = gr.Textbox(label="Antagonist", interactive=False)
                game_story_output = gr.Textbox(label="Game Story", interactive=False)
        
        submit_btn.click(generate_game_design, inputs=[environment, protagonist, antagonist], outputs=[env_output, prot_output, antag_output, game_story_output])
    
    demo.launch()

if __name__ == "__main__":
    main()