import pixeltable as pxt
import re

@pxt.udf
def generate_messages(genre: str, player_name: str, initial_scenario: str, player_input: str, turn_number: int) -> list[dict]:
    return [
        {
            'role': 'system',
            'content': f"""You are the game master for a {genre} RPG. The player's name is {player_name}. 
            Provide an engaging response to the player's action and present exactly 3 numbered options for their next move:
            1. A dialogue option (saying something)
            2. A random action they could take
            3. A unique or unexpected choice
            
            Format each option with a number (1., 2., 3.) followed by the action description."""
        },
        {
            'role': 'user',
            'content': f"Current scenario: {initial_scenario}\n"
                      f"Player's action: {player_input}\n"
                      f"Turn number: {turn_number}\n\n"
                      "Provide a response and 3 options:"
        }
    ]

@pxt.udf
def extract_options(story_text: str) -> list[str]:
    """Extract the three options from the story text"""
    # Look for numbered options (1., 2., 3.) and grab the text after them
    options = re.findall(r'\d\.\s*(.*?)(?=\d\.|$)', story_text, re.DOTALL)
    # Clean up the options and ensure we have exactly 3
    options = [opt.strip() for opt in options[:3]]
    while len(options) < 3:
        options.append("Take another action...")
    return options