import gradio as gr
import pixeltable as pxt
import numpy as np
from datetime import datetime
from pixeltable.functions.huggingface import sentence_transformer
from pixeltable.functions import openai
import os
import getpass
import re
import random
import time
import pandas as pd

# Set up OpenAI API key
if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('Enter your OpenAI API key: ')

# Initialize Pixeltable
pxt.drop_dir('ai_rpg', force=True)
pxt.create_dir('ai_rpg')

# Regular function (not UDF) for initializing stats
def initialize_stats(genre: str) -> str:
    """Initialize player stats based on the selected genre"""
    base_stats = {
        "🧙‍♂️ Fantasy": "Health: 100, Mana: 80, Strength: 7, Intelligence: 8, Agility: 6, Gold: 50",
        "🚀 Sci-Fi": "Health: 100, Energy: 90, Tech: 8, Intelligence: 9, Agility: 6, Credits: 500",
        "👻 Horror": "Health: 80, Sanity: 100, Strength: 6, Intelligence: 7, Agility: 8, Items: Flashlight, First Aid",
        "🔍 Mystery": "Health: 90, Focus: 100, Observation: 9, Intelligence: 8, Charisma: 7, Clues: 0",
        "🌋 Post-Apocalyptic": "Health: 95, Radiation Resistance: 75, Strength: 8, Survival: 9, Supplies: Limited",
        "🤖 Cyberpunk": "Health: 90, Cyberware: 85%, Hacking: 8, Street Cred: 6, Edge: 7, Nuyen: 1000",
        "⚙️ Steampunk": "Health: 95, Steam Power: 85, Engineering: 8, Artistry: 7, Social: 6, Shillings: 200"
    }
    
    if genre in base_stats:
        return base_stats[genre]
    else:
        # Default stats if genre not found
        return "Health: 100, Energy: 100, Strength: 7, Intelligence: 7, Agility: 7, Money: 100"

# Regular Python function for fallback responses (not a UDF)
def create_fallback_response(player_input: str, turn_number: int) -> dict:
    """Generate a fallback response when API fails"""
    if turn_number == 0:
        story = "Your adventure begins as you assess your surroundings. The world around you seems full of possibilities and dangers alike. What will you do next?"
        stats = "Your current stats remain unchanged."
        options = [
            "Explore the area carefully",
            "Look for signs of civilization",
            "Check your inventory",
            "Rest and gather your thoughts"
        ]
    else:
        story = f"You decide to {player_input}. As you do, you notice new details about your environment and feel yourself making progress in your quest."
        stats = "Your actions have slightly improved your experience."
        options = [
            "Continue on your current path",
            "Try a different approach",
            "Investigate something suspicious",
            "Take a moment to strategize"
        ]
    
    return {
        "story": story,
        "stats": stats,
        "options": options
    }

@pxt.udf
def generate_messages(genre: str, player_name: str, initial_scenario: str, player_input: str, turn_number: int, stats: str) -> list[dict]:
    return [
        {
            'role': 'system',
            'content': f"""You are the game master for a {genre} RPG. The player's name is {player_name}.
            
            Player stats to manage: {stats}
            
            You are a game master who vividly develops stories based on the player's choices.
            Use detailed descriptions and sensory details to help the player immerse themselves in the game world.
            
            When player stats change based on their choices, reflect this in the story.
            Create interesting stories that include dangerous situations, challenges, rewards, and chance encounters.
            
            Provide your response in three clearly separated sections using exactly this format:

            📜 **STORY**: [Your engaging narrative response to the player's action with vivid descriptions]

            📊 **STATS UPDATE**: [Brief update on any changes to player stats based on their actions]

            🎯 **OPTIONS**:
            1. [A dialogue option with potential consequences]
            2. [An action they could take with different outcomes]
            3. [A unique or unexpected choice that might lead to adventure]
            4. [A risky but potentially rewarding option]"""
        },
        {
            'role': 'user',
            'content': f"Current scenario: {initial_scenario}\n"
                      f"Player's action: {player_input}\n"
                      f"Turn number: {turn_number}\n"
                      f"Current player stats: {stats}\n\n"
                      "Provide the story response, stats update, and options:"
        }
    ]

@pxt.udf
def get_story(response: str) -> str:
    """Extract just the story part from the response"""
    match = re.search(r'📜\s*\*\*STORY\*\*:\s*(.*?)(?=📊\s*\*\*STATS|$)', response, re.DOTALL)
    if match:
        return match.group(1).strip()
    parts = response.split("STATS UPDATE:")
    if len(parts) > 1:
        story_part = parts[0].replace("STORY:", "").replace("📜", "").replace("**STORY**:", "").strip()
        return story_part
    return response

@pxt.udf
def get_stats_update(response: str) -> str:
    """Extract the stats update from the response"""
    match = re.search(r'📊\s*\*\*STATS UPDATE\*\*:\s*(.*?)(?=🎯\s*\*\*OPTIONS\*\*|$)', response, re.DOTALL)
    if match:
        return match.group(1).strip()
    parts = response.split("STATS UPDATE:")
    if len(parts) > 1:
        stats_part = parts[1].split("OPTIONS:")[0].strip()
        return stats_part
    return "No stat changes"

@pxt.udf
def get_options(response: str) -> list[str]:
    """Extract the options from the response"""
    match = re.search(r'🎯\s*\*\*OPTIONS\*\*:\s*(.*?)(?=$)', response, re.DOTALL)
    if match:
        options_text = match.group(1)
        options = re.findall(r'\d+\.\s*(.*?)(?=\d+\.|$)', options_text, re.DOTALL)
        options = [opt.strip() for opt in options if opt.strip()]
        while len(options) < 4:
            options.append("Try something else...")
        return options[:4]
    
    parts = response.split("OPTIONS:")
    if len(parts) > 1:
        options = re.findall(r'\d+\.\s*(.*?)(?=\d+\.|$)', parts[1], re.DOTALL)
        options = [opt.strip() for opt in options if opt.strip()]
        while len(options) < 4:
            options.append("Try something else...")
        return options[:4]
    
    return ["Continue...", "Take a different action", "Try something new", "Explore surroundings"]

# Create a single table for all game data
interactions = pxt.create_table(
    'ai_rpg.interactions',
    {
        'session_id': pxt.String,
        'player_name': pxt.String,
        'genre': pxt.String,
        'initial_scenario': pxt.String,
        'turn_number': pxt.Int,
        'player_input': pxt.String,
        'timestamp': pxt.Timestamp,
        'player_stats': pxt.String,
        'random_event': pxt.String,
        'use_fallback': pxt.Bool,
        'fallback_story': pxt.String,
        'fallback_stats': pxt.String,
        'fallback_options': pxt.String
    }
)

# Add computed columns for AI responses
interactions.add_computed_column(messages=generate_messages(
    interactions.genre,
    interactions.player_name,
    interactions.initial_scenario,
    interactions.player_input,
    interactions.turn_number,
    interactions.player_stats
))

# Changed to gpt-3.5-turbo for better compatibility
interactions.add_computed_column(ai_response=openai.chat_completions(
    messages=interactions.messages,
    model='gpt-3.5-turbo',
    max_tokens=800,
    temperature=0.8
))

interactions.add_computed_column(full_response=interactions.ai_response.choices[0].message.content)
interactions.add_computed_column(story_text=get_story(interactions.full_response))
interactions.add_computed_column(stats_update=get_stats_update(interactions.full_response))
interactions.add_computed_column(options=get_options(interactions.full_response))

class RPGGame:
    def __init__(self):
        self.current_session_id = None
        self.turn_number = 0
        self.current_stats = ""
        # Add game_history to track history directly in memory
        self.game_history = []
    
    def start_game(self, player_name: str, genre: str, scenario: str) -> tuple[str, str, list[str]]:
        session_id = f"session_{datetime.now().strftime('%Y%m%d%H%M%S')}_{player_name}"
        self.current_session_id = session_id
        self.turn_number = 0
        
        # Reset history for new game
        self.game_history = []
        
        # Get initial stats as a string
        initial_stats = initialize_stats(genre)
        self.current_stats = initial_stats
        
        # Create fallback content just in case
        fallback = create_fallback_response("Game starts", 0)
        
        try:
            # First try without fallback
            interactions.insert([{
                'session_id': session_id,
                'player_name': player_name,
                'genre': genre,
                'initial_scenario': scenario,
                'turn_number': 0,
                'player_input': "Game starts",
                'timestamp': datetime.now(),
                'player_stats': initial_stats,
                'random_event': "",
                'use_fallback': False,
                'fallback_story': fallback["story"],
                'fallback_stats': fallback["stats"],
                'fallback_options': ",".join(fallback["options"])
            }])
            
            # Try to get response
            try:
                result = interactions.select(
                    interactions.story_text,
                    interactions.stats_update,
                    interactions.options
                ).where(
                    (interactions.session_id == session_id) & 
                    (interactions.turn_number == 0)
                ).collect()
                
                story = result['story_text'][0]
                stats = result['stats_update'][0]
                options = result['options'][0]
                
                # Add to our internal history
                self.game_history.append(["0", "Game starts", story])
                
                return story, stats, options
            except Exception as e:
                print(f"Error getting initial response: {str(e)}")
                # If OpenAI fails, use the fallback values
                story = fallback["story"]
                stats = fallback["stats"]
                options = fallback["options"]
                
                # Still add to history even if using fallback
                self.game_history.append(["0", "Game starts", story])
                
                return story, stats, options
                
        except Exception as e:
            print(f"Error inserting initial turn: {str(e)}")
            # If everything fails, return fallback
            story = fallback["story"]
            stats = fallback["stats"]
            options = fallback["options"]
            
            # Still add to history even if using fallback
            self.game_history.append(["0", "Game starts", story])
            
            return story, stats, options
    
    def process_action(self, action: str) -> tuple[str, str, list[str]]:
        if not self.current_session_id:
            return "No active game session. Please start a new game.", "No stats", ["Start a new game"]
        
        self.turn_number += 1
        
        try:
            prev_turn = interactions.select(
                interactions.player_name,
                interactions.genre,
                interactions.initial_scenario,
                interactions.player_stats
            ).where(
                (interactions.session_id == self.current_session_id) &
                (interactions.turn_number == self.turn_number - 1)
            ).collect()
            
            self.current_stats = prev_turn['player_stats'][0]
            
            # Generate random event directly
            random_event_val = ""
            if self.turn_number % 3 == 0 and self.turn_number > 0:
                events = [
                    "Suddenly, you hear a strange sound nearby",
                    "A mysterious traveler is watching you",
                    "The ground begins to vibrate slightly",
                    "Something in your pocket starts to glow",
                    "Something is approaching you from the distance",
                    "The weather suddenly begins to change",
                    "You discover a hidden passage nearby"
                ]
                random_event_val = random.choice(events)
            
            modified_action = action
            if random_event_val:
                modified_action = f"{action} ({random_event_val})"
            
            # Create fallback content
            fallback = create_fallback_response(action, self.turn_number)
            
            interactions.insert([{
                'session_id': self.current_session_id,
                'player_name': prev_turn['player_name'][0],
                'genre': prev_turn['genre'][0],
                'initial_scenario': prev_turn['initial_scenario'][0],
                'turn_number': self.turn_number,
                'player_input': modified_action,
                'timestamp': datetime.now(),
                'player_stats': self.current_stats,
                'random_event': random_event_val,
                'use_fallback': False,
                'fallback_story': fallback["story"],
                'fallback_stats': fallback["stats"],
                'fallback_options': ",".join(fallback["options"])
            }])
            
            try:
                result = interactions.select(
                    interactions.story_text,
                    interactions.stats_update,
                    interactions.options
                ).where(
                    (interactions.session_id == self.current_session_id) & 
                    (interactions.turn_number == self.turn_number)
                ).collect()
                
                # Update stats for next turn
                story = result['story_text'][0]
                stats = result['stats_update'][0]
                options = result['options'][0]
                
                self.current_stats = stats
                
                # Add to our internal history
                self.game_history.append([str(self.turn_number), modified_action, story])
                
                return story, stats, options
            except Exception as e:
                print(f"Error getting turn response: {str(e)}")
                # If OpenAI fails, use the fallback values
                story = fallback["story"]
                stats = fallback["stats"]
                options = fallback["options"]
                
                # Add to our internal history even with fallback
                self.game_history.append([str(self.turn_number), modified_action, story])
                
                return story, stats, options
                
        except Exception as e:
            print(f"Error processing action: {str(e)}")
            # If everything fails, use fallback
            fallback = create_fallback_response(action, self.turn_number)
            story = fallback["story"]
            stats = fallback["stats"]
            options = fallback["options"]
            
            # Add to our internal history even with fallback
            self.game_history.append([str(self.turn_number), action, story])
            
            return story, stats, options
            
    def get_history(self):
        """Return the game history from memory instead of querying the database"""
        return self.game_history

def create_interface():
    game = RPGGame()
    
    # Custom CSS for improved visuals
    custom_css = """
    .container {
        max-width: 1200px;
        margin: 0 auto;
    }
    
    .title-container {
        background: linear-gradient(135deg, #6e48aa 0%, #9c27b0 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin-bottom: 20px;
        text-align: center;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    }
    
    .story-container {
        background: #f8f9fa;
        border-left: 5px solid #9c27b0;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-bottom: 20px;
        font-family: 'Roboto', sans-serif;
    }
    
    .stats-container {
        background: #e8f5e9;
        border-left: 5px solid #4caf50;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    
    .options-container {
        background: #e3f2fd;
        border-left: 5px solid #2196f3;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    
    .action-button {
        background: linear-gradient(135deg, #6e48aa 0%, #9c27b0 100%);
        color: white;
        border: none;
        padding: 10px 20px;
        border-radius: 5px;
        cursor: pointer;
        transition: all 0.3s ease;
    }
    
    .action-button:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 10px rgba(0,0,0,0.2);
    }
    
    .history-container {
        background: #fff8e1;
        border-left: 5px solid #ffc107;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-top: 20px;
    }
    """
    
    with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
        gr.HTML(
            """
            <div class="title-container">
                <h1 style="margin-bottom: 0.5em; font-size: 2.5em;">🎲 AI RPG Adventure</h1>
                <p style="font-size: 1.2em;">An immersive roleplaying experience powered by Pixeltable and OpenAI!</p>
            </div>
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Accordion("🎯 What is this app?", open=False):
                    gr.HTML(
                        """
                        <div style="padding: 15px;">
                            <h3>AI RPG Adventure showcases:</h3>
                            <ul style="list-style-type: none; padding-left: 5px;">
                                <li>🎮 <b>Dynamic Storytelling:</b> AI-generated immersive stories</li>
                                <li>🔄 <b>Game State Management:</b> Tracking game state with Pixeltable</li>
                                <li>💭 <b>Contextual Choices:</b> Custom options based on player actions</li>
                                <li>🤖 <b>AI Narratives:</b> Vivid storytelling</li>
                                <li>📊 <b>Character Status Tracking:</b> Stat changes as you play</li>
                            </ul>
                        </div>
                        """
                    )
                
                with gr.Accordion("🎨 Create Character", open=True):
                    player_name = gr.Textbox(
                        label="👤 Character Name",
                        placeholder="Enter your character's name...",
                        container=False
                    )
                    genre = gr.Dropdown(
                        choices=[
                            "🧙‍♂️ Fantasy",
                            "🚀 Sci-Fi",
                            "👻 Horror",
                            "🔍 Mystery",
                            "🌋 Post-Apocalyptic",
                            "🤖 Cyberpunk",
                            "⚙️ Steampunk"
                        ],
                        label="🎭 Choose Genre",
                        container=False,
                        value="🧙‍♂️ Fantasy"
                    )
                    scenario = gr.Textbox(
                        label="📖 Starting Scenario",
                        lines=3,
                        placeholder="Describe the initial setting and situation...",
                        container=False
                    )
                    start_button = gr.Button("🎮 Begin Adventure!", variant="primary")
            
            with gr.Column(scale=2):
                story_display = gr.Markdown(
                    label="📜 Story",
                    value="<div class='story-container'>Create a character and click 'Begin Adventure!' to start your journey.</div>",
                    show_label=False
                )
                
                stats_display = gr.Markdown(
                    label="📊 Character Stats",
                    value="<div class='stats-container'>Character stats will appear here when you start your adventure.</div>",
                    show_label=False
                )
                
                gr.HTML("<div class='options-container'><h3>🎯 Choose Your Next Action</h3></div>")
                
                action_input = gr.Radio(
                    choices=[],
                    label="",
                    interactive=True
                )
                submit_action = gr.Button("⚡ Take Action", variant="primary")
        
        with gr.Row():
            with gr.Column():
                gr.HTML("<div class='history-container'><h3>💫 Adventure Examples</h3></div>")
                gr.Examples(
                    examples=[
                        ["Admiral Yi", "🧙‍♂️ Fantasy", "You wake up at the edge of a forgotten mystical forest. In the distance, you can see castle spires, and in your mind remains only the clue about ancient magic threatening the kingdom. Suddenly, you see a strange light coming from the forest..."],
                        ["Ji-Young Kim", "🚀 Sci-Fi", "As the navigator of the starship 'Horizon', you wake up to an emergency alarm during the exploration of an unknown planet. Contact with the captain has been lost, and life support systems are gradually failing. You hear strange footsteps in the quiet ship..."],
                        ["Elon Musk", "🤖 Cyberpunk", "Seoul, 2077. You are the CEO of Neural Link Industries. Your latest brain-computer interface technology has started giving users an unexpected ability - they can communicate with their houseplants. As you prepare for a major investor presentation, your AI assistant reports that test participants are organizing a mysterious 'plant revolution'..."],
                        ["Gordon Ramsay", "🌋 Post-Apocalyptic", "You are the last master chef in New Seoul, running an underground restaurant in the ruins of an old luxury hotel. Your signature dish requires a rare mushroom that only grows in dangerous radioactive zones. While preparing for tonight's secret gathering, your scout returns with ominous news about rival chef gangs in the area..."],
                        ["Jung-Ho Yoo", "👻 Horror", "You came to a secluded cabin in the forest for the weekend at your friend's invitation. On the first night, you are drawn into the forest by a strange light visible through the window. As you try to find your way back, you can't see the cabin, and unfamiliar fog grows thicker. In the distance, you hear someone—no, something calling you..."],
                        ["Detective Park", "🔍 Mystery", "As a rookie detective, your first case is the mysterious disappearance of the CEO of the city's top tech company. There's no blood in his office, and the only clues are an encrypted note on the desk and his powered-down cutting-edge AI assistant. As you begin your investigation, you hear about a secret project the CEO was working on..."],
                        ["Min-Su Lee", "⚙️ Steampunk", "In New Joseon, filled with steam and gears, you are an innovative airship designer. During a demonstration of your latest invention, a government secret agent approaches with a secret mission for the imperial family in danger. Underground groups are threatening the throne, and your invention is the royal family's only hope..."]
                    ],
                    inputs=[player_name, genre, scenario]
                )
            
            with gr.Column():
                history_df = gr.Dataframe(
                    headers=["📅 Turn", "🎯 Player Action", "💬 Game Response"],
                    label="📚 Adventure History",
                    wrap=True,
                    row_count=5,
                    col_count=(3, "fixed")
                )

        def start_new_game(name, genre_choice, scenario_text):
            if not name or not genre_choice or not scenario_text:
                return (
                    "<div class='story-container'>Please fill in all fields before starting.</div>",
                    "<div class='stats-container'>No stat information</div>",
                    [],
                    []
                )
            
            try:
                initial_story, initial_stats, initial_options = game.start_game(name, genre_choice, scenario_text)
                
                # Get history from our in-memory tracking
                history_data = game.get_history()
                
                story_html = f"<div class='story-container'>{initial_story}</div>"
                stats_html = f"<div class='stats-container'><h3>📊 Character Status</h3>{initial_stats}</div>"
                
                return story_html, stats_html, gr.Radio(choices=initial_options, interactive=True), history_data
            except Exception as e:
                print(f"Error in start_new_game: {str(e)}")
                # Handle any other unexpected errors
                fallback = create_fallback_response("Game starts", 0)
                return (
                    f"<div class='story-container'>{fallback['story']}</div>",
                    f"<div class='stats-container'><h3>📊 Character Status</h3>{fallback['stats']}</div>",
                    gr.Radio(choices=fallback['options'], interactive=True),
                    [["0", "Game starts", fallback['story']]]
                )

        def process_player_action(action_choice):
            try:
                if not action_choice:
                    return (
                        "<div class='story-container'>Please select an action to continue.</div>",
                        "<div class='stats-container'>No stat information</div>",
                        [],
                        []
                    )
                    
                story, stats, options = game.process_action(action_choice)
                
                # Get history from our in-memory tracking
                history_data = game.get_history()
                
                story_html = f"<div class='story-container'>{story}</div>"
                stats_html = f"<div class='stats-container'><h3>📊 Character Status</h3>{stats}</div>"
                
                return story_html, stats_html, gr.Radio(choices=options, interactive=True), history_data
            except Exception as e:
                print(f"Error in process_player_action: {str(e)}")
                # Handle any other unexpected errors
                fallback = create_fallback_response(action_choice, 1)
                return (
                    f"<div class='story-container'>{fallback['story']}</div>",
                    f"<div class='stats-container'><h3>📊 Character Status</h3>{fallback['stats']}</div>",
                    gr.Radio(choices=fallback['options'], interactive=True),
                    [[str(game.turn_number), action_choice, fallback['story']]]
                )
        
        start_button.click(
            start_new_game,
            inputs=[player_name, genre, scenario],
            outputs=[story_display, stats_display, action_input, history_df]
        )
        
        submit_action.click(
            process_player_action,
            inputs=[action_input],
            outputs=[story_display, stats_display, action_input, history_df]
        )
        
        gr.HTML("""
        <div style="text-align: center; margin-top: 30px; padding: 20px; background: #f5f5f5; border-radius: 10px;">
            <h3>🌟 AI RPG Adventure - An Immersive Roleplaying Experience Powered by Pixeltable</h3>
            <p>Create your own character and embark on an adventure in a world of your chosen genre. Your choices shape the story!</p>
        </div>
        """)

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()