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

# 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)
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"

@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 fallback_response(player_input: str, turn_number: int) -> str:
    """Generate a fallback response when API fails"""
    if turn_number == 0:
        return """📜 **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 UPDATE**: Your current stats remain unchanged.

🎯 **OPTIONS**:
1. Explore the area carefully
2. Look for signs of civilization
3. Check your inventory
4. Rest and gather your thoughts"""
    else:
        return f"""📜 **STORY**: You decide to {player_input}. As you do, you notice new details about your environment and feel yourself making progress in your quest.

📊 **STATS UPDATE**: Your actions have slightly improved your experience.

🎯 **OPTIONS**:
1. Continue on your current path
2. Try a different approach
3. Investigate something suspicious
4. Take a moment to strategize"""

@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
    }
)

# 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 from gpt-4.1-mini 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',  # Changed to a more widely available model
    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 = ""
    
    def start_game(self, player_name: str, genre: str, scenario: str) -> tuple[str, 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
        
        # Get initial stats as a string
        initial_stats = initialize_stats(genre)
        self.current_stats = initial_stats
        
        try:
            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,  # Store string result
                'random_event': ""
            }])
            
            # Added retry logic for API calls
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    result = interactions.select(
                        interactions.story_text,
                        interactions.stats_update,
                        interactions.options
                    ).where(
                        (interactions.session_id == session_id) & 
                        (interactions.turn_number == 0)
                    ).collect()
                    
                    return session_id, result['story_text'][0], result['stats_update'][0], result['options'][0]
                except Exception as e:
                    if attempt < max_retries - 1:
                        time.sleep(2)  # Wait before retrying
                    else:
                        # Use fallback response if all retries fail
                        fallback = fallback_response("Game starts", 0)
                        return session_id, get_story(fallback), get_stats_update(fallback), get_options(fallback)
        except Exception as e:
            # Use fallback if insert fails
            fallback = fallback_response("Game starts", 0)
            return session_id, get_story(fallback), get_stats_update(fallback), get_options(fallback)
    
    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", []
        
        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)
            
            if random_event_val:
                action = f"{action} ({random_event_val})"
            
            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': action,
                'timestamp': datetime.now(),
                'player_stats': self.current_stats,
                'random_event': random_event_val
            }])
            
            # Added retry logic for API calls
            max_retries = 3
            for attempt in range(max_retries):
                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
                    self.current_stats = result['stats_update'][0]
                    
                    return result['story_text'][0], result['stats_update'][0], result['options'][0]
                except Exception as e:
                    if attempt < max_retries - 1:
                        time.sleep(2)  # Wait before retrying
                    else:
                        # Use fallback response if all retries fail
                        fallback = fallback_response(action, self.turn_number)
                        return get_story(fallback), get_stats_update(fallback), get_options(fallback)
        except Exception as e:
            # Use fallback if something fails
            fallback = fallback_response(action, self.turn_number)
            return get_story(fallback), get_stats_update(fallback), get_options(fallback)

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)
                
                history_df = interactions.select(
                    turn=interactions.turn_number,
                    action=interactions.player_input,
                    response=interactions.story_text
                ).where(
                    interactions.session_id == game.current_session_id
                ).order_by(
                    interactions.turn_number
                ).collect().to_pandas()
                
                history_data = [
                    [str(row['turn']), row['action'], row['response']] 
                    for _, row in history_df.iterrows()
                ]
                
                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:
                return (
                    f"<div class='story-container'>Game start error: {str(e)}</div>",
                    "<div class='stats-container'>No stat information</div>",
                    [],
                    []
                )

        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)
                
                history_df = interactions.select(
                    turn=interactions.turn_number,
                    action=interactions.player_input,
                    response=interactions.story_text
                ).where(
                    interactions.session_id == game.current_session_id
                ).order_by(
                    interactions.turn_number
                ).collect().to_pandas()
                
                history_data = [
                    [str(row['turn']), row['action'], row['response']] 
                    for _, row in history_df.iterrows()
                ]
                
                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:
                return (
                    f"<div class='story-container'>Error: {str(e)}</div>",
                    "<div class='stats-container'>No stat information</div>",
                    [],
                    []
                )
        
        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()