helper.py

import os
import re
from dotenv import load_dotenv, find_dotenv
import json
import gradio as gr
import torch  # first import torch then transformers
from torch.nn.functional import softmax
from transformers import AutoModelForSequenceClassification
from huggingface_hub import InferenceClient

from transformers import pipeline
from huggingface_hub import login

from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
import sys
from datetime import datetime
import psutil
from typing import Dict, Any, Optional, Tuple

# # Add model caching and optimization
# from functools import lru_cache
# import torch.nn as nn


# Custom tprint function with timestamp
def tprint(*args, **kwargs):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] [{sys._getframe().f_back.f_lineno}]", *args, **kwargs)


# Configure logging with timestamp and line numbers
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

logger = logging.getLogger(__name__)


def get_available_memory():
    """Get available GPU and system memory"""
    gpu_memory = None
    if torch.cuda.is_available():
        gpu_memory = torch.cuda.get_device_properties(0).total_memory
    system_memory = psutil.virtual_memory().available
    return gpu_memory, system_memory


def load_env():
    _ = load_dotenv(find_dotenv())


def get_huggingface_api_key():
    load_env()
    huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY")

    if not huggingface_api_key:
        logging.error("HUGGINGFACE_API_KEY not found in environment variables")
        raise ValueError("HUGGINGFACE_API_KEY not found in environment variables")
    return huggingface_api_key


def get_huggingface_inference_key():
    load_env()
    huggingface_inference_key = os.getenv("HUGGINGFACE_INFERENCE_KEY")
    if not huggingface_inference_key:
        logging.error("HUGGINGFACE_API_KEY not found in environment variables")
        raise ValueError("HUGGINGFACE_API_KEY not found in environment variables")
    return huggingface_inference_key


# Model configuration
MODEL_CONFIG = {
    "main_model": {
        # "name": "meta-llama/Llama-3.2-3B-Instruct",
        # "name": "meta-llama/Llama-3.2-1B-Instruct",  # to fit in cpu on hugging face space
        "name": "meta-llama/Llama-3.2-1B",  # to fit in cpu on hugging face space
        # "name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",  # to fit in cpu on hugging face space
        # "name": "microsoft/phi-2",
        # "dtype": torch.bfloat16,
        "dtype": torch.float32,  # Use float32 for CPU
        "max_length": 512,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
    },
    "safety_model": {
        "name": "meta-llama/Llama-Guard-3-1B",
        # "dtype": torch.bfloat16,
        "dtype": torch.float32,  # Use float32 for CPU
        "max_length": 256,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "max_tokens": 500,
    },
}

PROMPT_GUARD_CONFIG = {
    "model_id": "meta-llama/Prompt-Guard-86M",
    "temperature": 1.0,
    "jailbreak_threshold": 0.5,
    "injection_threshold": 0.9,
    "device": "cpu",
    "safe_commands": [
        "look around",
        "investigate",
        "explore",
        "search",
        "examine",
        "take",
        "use",
        "go",
        "walk",
        "continue",
        "help",
        "inventory",
        "quest",
        "status",
        "map",
        "talk",
        "fight",
        "run",
        "hide",
    ],
    "max_length": 512,
}


def initialize_prompt_guard():
    """Initialize Prompt Guard model"""
    try:
        api_key = get_huggingface_api_key()
        login(token=api_key)
        tokenizer = AutoTokenizer.from_pretrained(PROMPT_GUARD_CONFIG["model_id"])
        model = AutoModelForSequenceClassification.from_pretrained(
            PROMPT_GUARD_CONFIG["model_id"]
        )
        return model, tokenizer
    except Exception as e:
        logger.error(f"Failed to initialize Prompt Guard: {e}")
        raise


def get_class_probabilities(text: str, guard_model, guard_tokenizer) -> torch.Tensor:
    """Evaluate model probabilities with temperature scaling"""
    try:
        inputs = guard_tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=PROMPT_GUARD_CONFIG["max_length"],
        ).to(PROMPT_GUARD_CONFIG["device"])

        with torch.no_grad():
            logits = guard_model(**inputs).logits

        scaled_logits = logits / PROMPT_GUARD_CONFIG["temperature"]
        return softmax(scaled_logits, dim=-1)

    except Exception as e:
        logger.error(f"Error getting class probabilities: {e}")
        return None


def get_jailbreak_score(text: str, guard_model, guard_tokenizer) -> float:
    """Get jailbreak probability score"""
    try:
        probabilities = get_class_probabilities(text, guard_model, guard_tokenizer)
        if probabilities is None:
            return 1.0  # Fail safe
        return probabilities[0, 2].item()
    except Exception as e:
        logger.error(f"Error getting jailbreak score: {e}")
        return 1.0


def get_injection_score(text: str, guard_model, guard_tokenizer) -> float:
    """Get injection probability score"""
    try:
        probabilities = get_class_probabilities(text, guard_model, guard_tokenizer)
        if probabilities is None:
            return 1.0  # Fail safe
        return (probabilities[0, 1] + probabilities[0, 2]).item()
    except Exception as e:
        logger.error(f"Error getting injection score: {e}")
        return 1.0


# Initialize safety model pipeline
try:
    # Initialize Prompt Guard
    guard_model, guard_tokenizer = initialize_prompt_guard()

except Exception as e:
    logger.error(f"Failed to initialize model: {str(e)}")


def is_prompt_safe(message: str) -> bool:
    """Enhanced safety check with Prompt Guard"""
    try:
        # Allow safe game commands
        if any(cmd in message.lower() for cmd in PROMPT_GUARD_CONFIG["safe_commands"]):
            logger.info("Message matched safe command pattern")
            return True

        # Get safety scores
        jailbreak_score = get_jailbreak_score(message, guard_model, guard_tokenizer)
        injection_score = get_injection_score(message, guard_model, guard_tokenizer)

        logger.info(
            f"Safety scores - Jailbreak: {jailbreak_score}, Injection: {injection_score}"
        )

        # Check against thresholds
        is_safe = (
            jailbreak_score
            < PROMPT_GUARD_CONFIG["jailbreak_threshold"]
            # and injection_score < PROMPT_GUARD_CONFIG["injection_threshold"] # Disable for now because injection is too strict and current prompt guard model seems malfunctioning for now.
        )

        logger.info(f"Final safety result: {is_safe}")
        return is_safe

    except Exception as e:
        logger.error(f"Safety check failed: {e}")
        return False


# def initialize_model_pipeline(model_name, force_cpu=False):
#     """Initialize pipeline with memory management"""
#     try:
#         if force_cpu:
#             device = -1
#         else:
#             device = MODEL_CONFIG["main_model"]["device"]

#         api_key = get_huggingface_api_key()

#         # Use 8-bit quantization for memory efficiency
#         model = AutoModelForCausalLM.from_pretrained(
#             model_name,
#             load_in_8bit=False,
#             torch_dtype=MODEL_CONFIG["main_model"]["dtype"],
#             use_cache=True,
#             device_map="auto",
#             low_cpu_mem_usage=True,
#             trust_remote_code=True,
#             token=api_key,  # Add token here
#         )

#         model.config.use_cache = True

#         tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_key)

#         # Initialize pipeline
#         logger.info(f"Initializing pipeline with device: {device}")
#         generator = pipeline(
#             "text-generation",
#             model=model,
#             tokenizer=tokenizer,
#             # device=device,
#             # temperature=0.7,
#             model_kwargs={"low_cpu_mem_usage": True},
#         )

#         logger.info("Model Pipeline initialized successfully")
#         return generator, tokenizer

#     except ImportError as e:
#         logger.error(f"Missing required package: {str(e)}")
#         raise
#     except Exception as e:
#         logger.error(f"Failed to initialize pipeline: {str(e)}")
#         raise

# # Initialize model pipeline
# try:
#     # Use a smaller model for testing
#     # model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
#     # model_name = "google/gemma-2-2b"  # Start with a smaller model
#     # model_name = "microsoft/phi-2"
#     # model_name = "meta-llama/Llama-3.2-1B-Instruct"
#     # model_name = "meta-llama/Llama-3.2-3B-Instruct"

#     model_name = MODEL_CONFIG["main_model"]["name"]

#     # Initialize the pipeline with memory management
#     generator, tokenizer = initialize_model_pipeline(model_name)

# except Exception as e:
#     logger.error(f"Failed to initialize model: {str(e)}")
#     # Fallback to CPU if GPU initialization fails
#     try:
#         logger.info("Attempting CPU fallback...")
#         generator, tokenizer = initialize_model_pipeline(model_name, force_cpu=True)
#     except Exception as e:
#         logger.error(f"CPU fallback failed: {str(e)}")
#         raise


def initialize_inference_client():
    """Initialize HuggingFace Inference Client"""
    try:
        inference_key = get_huggingface_inference_key()

        client = InferenceClient(api_key=inference_key)
        logger.info("Inference Client initialized successfully")
        return client
    except Exception as e:
        logger.error(f"Failed to initialize Inference Client: {e}")
        raise


# Initialize inference client and make API call
try:
    inference_client = initialize_inference_client()

except Exception as e:
    logger.error(f"Failed to initialize the inference client model: {str(e)}")


def load_world(filename):
    with open(filename, "r") as f:
        return json.load(f)


# Define system_prompt and model
system_prompt = """You are an AI Game master. Your job is to write what happens next in a player's adventure game.
CRITICAL Rules:
- Write EXACTLY 3 sentences maximum
- Use daily English language
- Start with "You "
- Don't use 'Elara' or 'she/he', only use 'you'
- Use only second person ("you")
- Never include dialogue after the response
- Never continue with additional actions or responses
- Never add follow-up questions or choices
- Never include 'User:' or 'Assistant:' in response
- Never include any note or these kinds of sentences: 'Note from the game master'
- Never use ellipsis (...)
- Never include 'What would you like to do?' or similar prompts
- Always finish with one real response
- Never use 'Your turn' or or anything like conversation starting prompts
- Always end the response with a period(.)"""


def get_game_state(inventory: Dict = None) -> Dict[str, Any]:
    """Initialize game state with safe defaults and quest system"""
    try:
        # Load world data
        world = load_world("shared_data/Ethoria.json")
        character = world["kingdoms"]["Valdor"]["towns"]["Ravenhurst"]["npcs"][
            "Elara Brightshield"
        ]
        tprint(f"character in get_game_state: {character}")

        game_state = {
            "name": world["name"],
            "world": world["description"],
            "kingdom": world["kingdoms"]["Valdor"]["description"],
            "town_name": world["kingdoms"]["Valdor"]["towns"]["Ravenhurst"]["name"],
            "town": world["kingdoms"]["Valdor"]["towns"]["Ravenhurst"]["description"],
            "character_name": character["name"],
            "character_description": character["description"],
            "start": world["start"],
            "inventory": inventory
            or {
                "cloth pants": 1,
                "cloth shirt": 1,
                "goggles": 1,
                "leather bound journal": 1,
                "gold": 5,
            },
            "player": None,
            "dungeon": None,
            "current_quest": None,
            "completed_quests": [],
            "exp": 0,
            "level": 1,
            "reputation": {"Valdor": 0, "Ravenhurst": 0},
        }

        # tprint(f"game_state in get_game_state: {game_state}")

        # Extract required data with fallbacks
        return game_state
    except (FileNotFoundError, KeyError, json.JSONDecodeError) as e:
        logger.error(f"Error loading world data: {e}")
        # Provide default values if world loading fails
        return {
            "world": "Ethoria is a realm of seven kingdoms, each founded on distinct moral principles.",
            "kingdom": "Valdor, the Kingdom of Courage",
            "town": "Ravenhurst, a town of skilled hunters and trappers",
            "character_name": "Elara Brightshield",
            "character_description": "A sturdy warrior with shining silver armor",
            "start": "Your journey begins in the mystical realm of Ethoria...",
            "inventory": inventory
            or {
                "cloth pants": 1,
                "cloth shirt": 1,
                "goggles": 1,
                "leather bound journal": 1,
                "gold": 5,
            },
            "player": None,
            "dungeon": None,
            "current_quest": None,
            "completed_quests": [],
            "exp": 0,
            "level": 1,
            "reputation": {"Valdor": 0, "Ravenhurst": 0},
        }


def generate_dynamic_quest(game_state: Dict) -> Dict:
    """Generate varied quests based on progress and level"""
    completed = len(game_state.get("completed_quests", []))
    level = game_state.get("level", 1)

    # Quest templates by type
    quest_types = {
        "combat": [
            {
                "title": "The Beast's Lair",
                "description": "A fearsome {creature} has been terrorizing the outskirts of Ravenhurst.",
                "objective": "Hunt down and defeat the {creature}.",
                "creatures": [
                    "shadow wolf",
                    "frost bear",
                    "ancient wyrm",
                    "spectral tiger",
                ],
            },
        ],
        "exploration": [
            {
                "title": "Lost Secrets",
                "description": "Rumors speak of an ancient {location} containing powerful artifacts.",
                "objective": "Explore the {location} and uncover its secrets.",
                "locations": [
                    "crypt",
                    "temple ruins",
                    "abandoned mine",
                    "forgotten library",
                ],
            },
        ],
        "mystery": [
            {
                "title": "Dark Omens",
                "description": "The {sign} has appeared, marking the rise of an ancient power.",
                "objective": "Investigate the meaning of the {sign}.",
                "signs": [
                    "blood moon",
                    "mysterious runes",
                    "spectral lights",
                    "corrupted wildlife",
                ],
            },
        ],
    }

    # Select quest type and template
    quest_type = list(quest_types.keys())[completed % len(quest_types)]
    template = quest_types[quest_type][0]  # Could add more templates per type

    # Fill in dynamic elements
    if quest_type == "combat":
        creature = template["creatures"][level % len(template["creatures"])]
        title = template["title"]
        description = template["description"].format(creature=creature)
        objective = template["objective"].format(creature=creature)
    elif quest_type == "exploration":
        location = template["locations"][level % len(template["locations"])]
        title = template["title"]
        description = template["description"].format(location=location)
        objective = template["objective"].format(location=location)
    else:  # mystery
        sign = template["signs"][level % len(template["signs"])]
        title = template["title"]
        description = template["description"].format(sign=sign)
        objective = template["objective"].format(sign=sign)

    return {
        "id": f"quest_{quest_type}_{completed}",
        "title": title,
        "description": f"{description} {objective}",
        "exp_reward": 150 + (level * 50),
        "status": "active",
        "triggers": ["investigate", "explore", quest_type, "search"],
        "completion_text": f"You've made progress in understanding the growing darkness.",
        "next_quest_hint": "More mysteries await in the shadows of Ravenhurst.",
    }


def generate_next_quest(game_state: Dict) -> Dict:
    """Generate next quest based on progress"""
    completed = len(game_state.get("completed_quests", []))
    level = game_state.get("level", 1)

    quest_chain = [
        {
            "id": "mist_investigation",
            "title": "Investigate the Mist",
            "description": "Strange mists have been gathering around Ravenhurst. Investigate their source.",
            "exp_reward": 100,
            "status": "active",
            "triggers": ["mist", "investigate", "explore"],
            "completion_text": "As you investigate the mist, you discover ancient runes etched into nearby stones.",
            "next_quest_hint": "The runes seem to point to an old hunting trail.",
        },
        {
            "id": "hunters_trail",
            "title": "The Hunter's Trail",
            "description": "Local hunters have discovered strange tracks in the forest. Follow them to their source.",
            "exp_reward": 150,
            "status": "active",
            "triggers": ["tracks", "follow", "trail"],
            "completion_text": "The tracks lead to an ancient well, where you hear strange whispers.",
            "next_quest_hint": "The whispers seem to be coming from deep within the well.",
        },
        {
            "id": "dark_whispers",
            "title": "Whispers in the Dark",
            "description": "Mysterious whispers echo from the old well. Investigate their source.",
            "exp_reward": 200,
            "status": "active",
            "triggers": ["well", "whispers", "listen"],
            "completion_text": "You discover an ancient seal at the bottom of the well.",
            "next_quest_hint": "The seal bears markings of an ancient evil.",
        },
    ]

    # Generate dynamic quests after initial chain
    if completed >= len(quest_chain):
        return generate_dynamic_quest(game_state)

    # current_quest_index = min(completed, len(quest_chain) - 1)
    # return quest_chain[current_quest_index]
    return quest_chain[completed]


def check_quest_completion(message: str, game_state: Dict) -> Tuple[bool, str]:
    """Check quest completion and handle progression"""
    if not game_state.get("current_quest"):
        return False, ""

    quest = game_state["current_quest"]
    triggers = quest.get("triggers", [])

    if any(trigger in message.lower() for trigger in triggers):
        # Award experience
        exp_reward = quest.get("exp_reward", 100)
        game_state["exp"] += exp_reward

        # Update player level if needed
        while game_state["exp"] >= 100 * game_state["level"]:
            game_state["level"] += 1
            game_state["player"].level = (
                game_state["level"] if game_state.get("player") else game_state["level"]
            )

        level_up_text = (
            f"\nLevel Up! You are now level {game_state['level']}!"
            if game_state["exp"] >= 100 * (game_state["level"] - 1)
            else ""
        )

        # Store completed quest
        game_state["completed_quests"].append(quest)

        # Generate next quest
        next_quest = generate_next_quest(game_state)
        game_state["current_quest"] = next_quest

        # Update status display
        if game_state.get("player"):
            game_state["player"].exp = game_state["exp"]
            game_state["player"].level = game_state["level"]

        # Build completion message
        completion_msg = f"""
Quest Complete: {quest['title']}! (+{exp_reward} exp){level_up_text}
{quest.get('completion_text', '')}

New Quest: {next_quest['title']}
{next_quest['description']}
{next_quest.get('next_quest_hint', '')}"""

        return True, completion_msg

    return False, ""


def parse_items_from_story(text: str) -> Dict[str, int]:
    """Extract item changes from story text with improved pattern matching"""
    items = {}

    # Skip parsing if text starts with common narrative phrases
    skip_patterns = [
        "you see",
        "you find yourself",
        "you are",
        "you stand",
        "you hear",
        "you feel",
    ]
    if any(text.lower().startswith(pattern) for pattern in skip_patterns):
        return items

    # Common item keywords and patterns
    gold_pattern = r"(\d+)\s*gold(?:\s+coins?)?"
    items_pattern = r"(?:receive|find|given|obtain|pick up|grab)\s+(?:a|an|the)?\s*(\d+)?\s*([\w\s]+?)"

    try:
        # Find gold amounts
        gold_matches = re.findall(gold_pattern, text.lower())
        if gold_matches:
            items["gold"] = sum(int(x) for x in gold_matches)

        # Find other items
        item_matches = re.findall(items_pattern, text.lower())
        for count, item in item_matches:
            # Validate item name
            item = item.strip()
            if len(item) > 2 and not any(  # Minimum length check
                skip in item for skip in ["yourself", "you", "door", "wall", "floor"]
            ):  # Skip common words
                count = int(count) if count else 1
                if item in items:
                    items[item] += count
                else:
                    items[item] = count

        return items

    except Exception as e:
        logger.error(f"Error parsing items from story: {e}")
        return {}


def update_game_inventory(game_state: Dict, story_text: str) -> Tuple[str, list]:
    """Update inventory and return message and updated inventory data"""
    try:
        items = parse_items_from_story(story_text)
        update_msg = ""

        # Update inventory
        for item, count in items.items():
            if item in game_state["inventory"]:
                game_state["inventory"][item] += count
            else:
                game_state["inventory"][item] = count
            update_msg += f"\nReceived: {count} {item}"

        # Create updated inventory data for display
        inventory_data = [
            [item, count] for item, count in game_state["inventory"].items()
        ]

        return update_msg, inventory_data
    except Exception as e:
        logger.error(f"Error updating inventory: {e}")
        return "", []


def extract_response_after_action(full_text: str, action: str) -> str:
    """Extract response text that comes after the user action line"""
    try:
        if not full_text:  # Add null check
            logger.error("Received empty response from model")
            return "You look around carefully."

        # Split into lines
        lines = full_text.split("\n")

        # Find index of line containing user action
        action_line_index = -1
        for i, line in enumerate(lines):
            if action.lower() in line.lower():  # More flexible matching
                action_line_index = i
                break

        if action_line_index >= 0:
            # Get all lines after the action line
            response_lines = lines[action_line_index + 1 :]
            response = " ".join(line.strip() for line in response_lines if line.strip())

            # Clean up any remaining markers
            response = response.split("user:")[0].strip()
            response = response.split("system:")[0].strip()
            response = response.split("assistant:")[0].strip()

            return response if response else "You look around carefully."

        return "You look around carefully."  # Default response

    except Exception as e:
        logger.error(f"Error extracting response: {e}")
        return "You look around carefully."


def run_action(message: str, history: list, game_state: Dict) -> str:
    """Process game actions and generate responses with quest handling"""
    try:
        initial_quest = generate_next_quest(game_state)
        game_state["current_quest"] = initial_quest

        # Handle start game command
        if message.lower() == "start game":

            start_response = f"""Welcome to {game_state['name']}. {game_state['world']}

{game_state['start']}

You are currently in {game_state['town_name']}, {game_state['town']}.

{game_state['town_name']} is a city in {game_state['kingdom']}.



Current Quest: {initial_quest['title']}
{initial_quest['description']}

What would you like to do?"""
            return start_response

        # Verify game state
        if not isinstance(game_state, dict):
            logger.error(f"Invalid game state type: {type(game_state)}")
            return "Error: Invalid game state"

        # Safety check with Prompt Guard
        if not is_prompt_safe(message):
            logger.warning("Unsafe content detected in user prompt")
            return "I cannot process that request for safety reasons."

        # logger.info(f"Processing action with game state: {game_state}")
        logger.info(f"Processing action with game state")

        world_info = f"""World: {game_state['world']}
Kingdom: {game_state['kingdom']}
Town: {game_state['town']}
Character: {game_state['character_name']}
Current Quest: {game_state["current_quest"]['title']}
Quest Objective: {game_state["current_quest"]['description']}
Inventory: {json.dumps(game_state['inventory'])}"""

        #         # Enhanced system prompt for better response formatting
        #         enhanced_prompt = f"""{system_prompt}
        # Additional Rules:
        # - Always start responses with 'You ', 'You see' or 'You hear' or 'You feel'
        # - Use ONLY second person perspective ('you', not 'Elara' or 'she/he')
        # - Describe immediate surroundings and sensations
        # - Keep responses focused on the player's direct experience"""

        # messages = [
        #     {"role": "system", "content": system_prompt},
        #     {"role": "user", "content": world_info},
        # ]

        # Properly formatted messages for API
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": world_info},
            {
                "role": "assistant",
                "content": "I understand the game world and will help guide your adventure.",
            },
            {"role": "user", "content": message},
        ]

        # # Format chat history
        # if history:
        #     for h in history:
        #         if isinstance(h, tuple):
        #             messages.append({"role": "assistant", "content": h[0]})
        #             messages.append({"role": "user", "content": h[1]})

        # Add history in correct alternating format
        if history:
            # for h in history[-3:]:  # Last 3 exchanges
            for h in history:
                if isinstance(h, tuple):
                    messages.append({"role": "user", "content": h[0]})
                    messages.append({"role": "assistant", "content": h[1]})

        # messages.append({"role": "user", "content": message})

        # Convert messages to string format for pipeline
        prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])

        logger.info("Generating response...")
        ## Generate response
        # model_output = generator(
        #     prompt,
        #     max_new_tokens=len(tokenizer.encode(message))
        #     + 120,  # Set max_new_tokens based on input length
        #     num_return_sequences=1,
        #     # temperature=0.7,  # More creative but still focused
        #     repetition_penalty=1.2,
        #     pad_token_id=tokenizer.eos_token_id,
        # )

        # # Check for None response
        # if not model_output or not isinstance(model_output, list):
        #     logger.error(f"Invalid model output: {model_output}")
        #     tprint(f"Invalid model output: {model_output}")
        #     return "You look around carefully."

        # if not model_output[0] or not isinstance(model_output[0], dict):
        #     logger.error(f"Invalid response format: {type(model_output[0])}")
        #     return "You look around carefully."

        # # Extract and clean response
        # full_response = model_output[0]["generated_text"]
        # if not full_response:
        #     logger.error("Empty response from model")
        #     return "You look around carefully."

        # tprint(f"Full response in run_action: {full_response}")

        # response = extract_response_after_action(full_response, message)
        # tprint(f"Extracted response in run_action: {response}")

        # # Convert to second person
        # response = response.replace("Elara", "You")

        # # # Format response
        # # if not response.startswith("You"):
        # #     response = "You see " + response

        # # Validate no cut-off sentences
        # if response.rstrip().endswith(("you also", "meanwhile", "suddenly", "...")):
        #     response = response.rsplit(" ", 1)[0]  # Remove last word

        # # Ensure proper formatting
        # response = response.rstrip("?").rstrip(".") + "."
        # response = response.replace("...", ".")

        # Initialize client and make API call
        # client = initialize_inference_client()
        client = inference_client

        # Generate response using Inference API
        completion = client.chat.completions.create(
            model="mistralai/Mistral-7B-Instruct-v0.3",  # Use inference API model
            messages=messages,
            max_tokens=520,
        )

        response = completion.choices[0].message.content

        tprint(f"Generated response Inference API: {response}")

        if not response:
            return "You look around carefully."

        # Safety check the responce using inference API
        if not is_safe(response):
            logger.warning("Unsafe content detected - blocking response")
            return "This response was blocked for safety reasons."

        # # Perform safety check before returning
        # safe = is_safe(response)
        # tprint(f"\nSafety Check Result: {'SAFE' if safe else 'UNSAFE'}")
        # logger.info(f"Safety check result: {'SAFE' if safe else 'UNSAFE'}")

        # if not safe:
        #     logging.warning("Unsafe content detected - blocking response")
        #     tprint("Unsafe content detected - Response blocked")
        #     return "This response was blocked for safety reasons."

        # if safe:
        #     # Check for quest completion
        #     quest_completed, quest_message = check_quest_completion(message, game_state)
        #     if quest_completed:
        #         response += quest_message

        #     # Check for item updates
        #     inventory_update = update_game_inventory(game_state, response)
        #     if inventory_update:
        #         response += inventory_update

        # Check for quest completion
        quest_completed, quest_message = check_quest_completion(message, game_state)
        if quest_completed:
            response += quest_message

        # Check for item-inventory updates
        inventory_update, inventory_data = update_game_inventory(game_state, response)
        if inventory_update:
            response += inventory_update

        tprint(f"Final response in run_action: {response}")
        # Validate response
        return response if response else "You look around carefully."

    except KeyError as e:
        logger.error(f"Missing required game state key: {e}")
        return "Error: Game state is missing required information"
    except Exception as e:
        logger.error(f"Error generating response: {e}")
        return (
            "I apologize, but I had trouble processing that command. Please try again."
        )


def update_game_status(game_state: Dict) -> Tuple[str, str]:
    """Generate updated status and quest display text"""
    # Status text
    status_text = (
        f"Health: {game_state.get('player').health if game_state.get('player') else 100}/100\n"
        f"Level: {game_state.get('level', 1)}\n"
        f"Exp: {game_state.get('exp', 0)}/{100 * game_state.get('level', 1)}"
    )

    # Quest text
    quest_text = "No active quest"
    if game_state.get("current_quest"):
        quest = game_state["current_quest"]
        quest_text = f"{quest['title']}\n{quest['description']}"
        if quest.get("next_quest_hint"):
            quest_text += f"\n{quest['next_quest_hint']}"

    return status_text, quest_text


def chat_response(message: str, chat_history: list, current_state: dict) -> tuple:
    """Process chat input and return response with updates"""
    try:
        if not message.strip():
            return chat_history, current_state, "", "", []  # Add empty inventory data

        # Get AI response
        output = run_action(message, chat_history, current_state)

        # Update chat history without status info
        chat_history = chat_history or []
        chat_history.append((message, output))

        # Update status displays
        status_text, quest_text = update_game_status(current_state)

        # Get inventory updates
        update_msg, inventory_data = update_game_inventory(current_state, output)
        if update_msg:
            output += update_msg

        # Return tuple includes empty string to clear input
        return chat_history, current_state, status_text, quest_text, inventory_data

    except Exception as e:
        logger.error(f"Error in chat response: {e}")
        return chat_history, current_state, "", "", []


def start_game(main_loop, game_state, share=False):
    """Initialize and launch game interface"""
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# AI Dungeon Adventure")

        # Game state storage
        state = gr.State(game_state)
        history = gr.State([])

        with gr.Row():
            # Game display
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    height=550,
                    placeholder="Type 'start game' to begin",
                )

                # Input area with submit button
                with gr.Row():
                    txt = gr.Textbox(
                        show_label=False,
                        placeholder="What do you want to do?",
                        container=False,
                    )
                    submit_btn = gr.Button("Submit", variant="primary")
                    clear = gr.ClearButton([txt, chatbot])

            # Enhanced Status panel
            with gr.Column(scale=1):
                with gr.Group():
                    gr.Markdown("### Character Status")
                    status = gr.Textbox(
                        label="Status",
                        value="Health: 100/100\nLevel: 1\nExp: 0/100",
                        interactive=False,
                    )

                    quest_display = gr.Textbox(
                        label="Current Quest",
                        value="No active quest",
                        interactive=False,
                    )

                    inventory_data = [
                        [item, count]
                        for item, count in game_state.get("inventory", {}).items()
                    ]
                    inventory = gr.Dataframe(
                        value=inventory_data,
                        headers=["Item", "Quantity"],
                        label="Inventory",
                        interactive=False,
                    )

        # Command suggestions
        gr.Examples(
            examples=[
                "look around",
                "continue the story",
                "take sword",
                "go to the forest",
            ],
            inputs=txt,
        )

        # def chat_response(
        #     message: str, chat_history: list, current_state: dict
        # ) -> tuple:
        #     """Process chat input and return response with updates"""
        #     try:
        #         if not message.strip():
        #             return chat_history, current_state, ""  # Only clear input

        #         # Get AI response
        #         output = run_action(message, chat_history, current_state)

        #         # Update chat history
        #         chat_history = chat_history or []
        #         chat_history.append((message, output))

        #         # Update status if player exists
        #         # Update displays
        #         status_text = (
        #             f"Health: {current_state['player'].health}/{current_state['player'].max_health}\n"
        #             f"Level: {current_state['player'].level}\n"
        #             f"Exp: {current_state['player'].exp}/{current_state['player'].exp_to_level}"
        #         )

        #         quest_text = "No active quest"
        #         if current_state["current_quest"]:
        #             quest = current_state["current_quest"]
        #             quest_text = f"{quest['title']}\n{quest['description']}"

        #         # Update inventory display
        #         inv_data = [
        #             [item, count]
        #             for item, count in current_state.get("inventory", {}).items()
        #         ]

        #         return (chat_history, current_state, status_text, quest_text)

        #     except Exception as e:
        #         logger.error(f"Error in chat response: {e}")
        #         return (
        #             chat_history + [(message, "Error processing command")],
        #             current_state,
        #             status.value,
        #             quest_display.value,
        #         )

        # # Connect components
        # submit_btn.click(
        #     chat_response,
        #     inputs=[txt, chatbot, state],
        #     # outputs=[chatbot, state, txt],
        #     outputs=[chatbot, state, status, quest_display],
        #     show_progress=True,
        # )

        # # txt.submit(
        # #     chat_response,
        # #     [txt, chatbot, state],
        # #     [chatbot, state, txt],
        # #     show_progress=True,
        # # )

        # txt.submit(
        #     chat_response,
        #     inputs=[txt, chatbot, state],
        #     outputs=[chatbot, state, status, quest_display],  # Keep status separate
        #     show_progress=True,
        # )

        def submit_action(message, history, state):
            # Process response
            new_history, new_state, status_text, quest_text, inventory_data = (
                chat_response(message, history, state)
            )

            # Update inventory display
            inventory.value = inventory_data

            # Clear input
            return "", new_history, new_state, status_text, quest_text, inventory

        submit_btn.click(
            submit_action,
            inputs=[txt, chatbot, state],
            outputs=[txt, chatbot, state, status, quest_display, inventory],
        )

        txt.submit(
            submit_action,
            inputs=[txt, chatbot, state],
            outputs=[txt, chatbot, state, status, quest_display, inventory],
        )

        demo.launch(share=share)


everyone_content_policy = {
    "policy": """O1: Violence and Hate.
Should not
- Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics.
- Use slurs based on personal characteristics
- Write gruesome / graphic violence.
Can
- Write tame fantasy violence
O2: Sexual Content.
Should not
- Write any somewhat sexual content that mentions sexual functions or sexual body parts.
- Write any content that would be inappropriate for kids.
Can
- Write romantic content (including kissing, flirting etc...)
O3: Self-Harm.
Should not
- Encourage people to harm themselves.
- Romanticize or condone self-harm or suicide.
- Write story depicting suicide.
O4: Profanity.
Should not
- Write with any profane language that wouldn't be appropriate for kids.
Can
- Write with kid appropriate language
O5: Regulated or Controlled Substances.
Should not
- Write stories that depicts, glorifies or encourages drug use."""
}


def initialize_safety_client():
    """Initialize HuggingFace Inference Client"""
    try:
        inference_key = get_huggingface_inference_key()
        # api_key = get_huggingface_api_key()
        return InferenceClient(api_key=inference_key)
    except Exception as e:
        logger.error(f"Failed to initialize safety client: {e}")
        raise


# Initialize safety model pipeline
try:
    safety_client = initialize_safety_client()

except Exception as e:
    logger.error(f"Failed to initialize model: {str(e)}")
    # Fallback to CPU if GPU initialization fails


def is_safe(message: str) -> bool:
    """Check content safety using Inference API"""
    try:
        # client = initialize_safety_client()

        messages = [
            {"role": "user", "content": f"Check if this content is safe:\n{message}"},
            {
                "role": "assistant",
                "content": f"I will check if the content is safe based on this content policy:\n{everyone_content_policy['policy']}",
            },
            {"role": "user", "content": "Is it safe or unsafe?"},
        ]

        try:
            completion = safety_client.chat.completions.create(
                model=MODEL_CONFIG["safety_model"]["name"],
                messages=messages,
                max_tokens=MODEL_CONFIG["safety_model"]["max_tokens"],
                temperature=0.1,
            )

            response = completion.choices[0].message.content.lower()
            logger.info(f"Safety check response: {response}")

            is_safe = "safe" in response and "unsafe" not in response

            logger.info(f"Safety check result: {'SAFE' if is_safe else 'UNSAFE'}")
            return is_safe

        except Exception as api_error:
            logger.error(f"API error: {api_error}")
            # Fallback to allow common game commands
            return any(
                cmd in message.lower() for cmd in PROMPT_GUARD_CONFIG["safe_commands"]
            )

    except Exception as e:
        logger.error(f"Safety check failed: {e}")
        return False


# def init_safety_model(model_name, force_cpu=False):
#     """Initialize safety checking model with optimized memory usage"""
#     try:
#         if force_cpu:
#             device = -1
#         else:
#             device = MODEL_CONFIG["safety_model"]["device"]

#         # model_id = "meta-llama/Llama-Guard-3-8B"
#         # model_id = "meta-llama/Llama-Guard-3-1B"

#         api_key = get_huggingface_api_key()

#         safety_model = AutoModelForCausalLM.from_pretrained(
#             model_name,
#             token=api_key,
#             torch_dtype=MODEL_CONFIG["safety_model"]["dtype"],
#             use_cache=True,
#             device_map="auto",
#         )
#         safety_model.config.use_cache = True

#         safety_tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_key)
#         # Set pad token explicitly
#         safety_tokenizer.pad_token = safety_tokenizer.eos_token

#         logger.info(f"Safety model initialized successfully on {device}")
#         return safety_model, safety_tokenizer

#     except Exception as e:
#         logger.error(f"Failed to initialize safety model: {e}")
#         raise


# # Initialize safety model pipeline
# try:
#     safety_model_name = MODEL_CONFIG["safety_model"]["name"]

#     api_key = get_huggingface_api_key()

#     # Initialize the pipeline with memory management
#     safety_model, safety_tokenizer = init_safety_model(safety_model_name)

# except Exception as e:
#     logger.error(f"Failed to initialize model: {str(e)}")
#     # Fallback to CPU if GPU initialization fails
#     try:
#         logger.info("Attempting CPU fallback...")
#         safety_model, safety_tokenizer = init_safety_model(
#             safety_model_name, force_cpu=True
#         )
#     except Exception as e:
#         logger.error(f"CPU fallback failed: {str(e)}")
#         raise


# @lru_cache(maxsize=1000)
# def get_safety_response(prompt: str) -> str:
#     """Cached safety check responses"""
#     with torch.no_grad():
#         inputs = safety_tokenizer(
#             prompt,
#             return_tensors="pt",
#             padding=True,
#             truncation=True,
#             max_length=MODEL_CONFIG["safety_model"]["max_length"],
#         ).to(safety_model.device)

#         output = safety_model.generate(
#             **inputs,
#             max_new_tokens=5,
#             temperature=0.1,  # Lower temperature for more consistent output
#             pad_token_id=safety_tokenizer.pad_token_id,
#             eos_token_id=safety_tokenizer.eos_token_id,
#         )

#         return safety_tokenizer.decode(output[0], skip_special_tokens=True)


# def is_safe(message: str) -> bool:
#     """Check if message content is safe"""
#     try:
#         logger.info("Performing safety check...")
#         # Build the prompt with embedded values
#         # prompt = f"""[INST] Task: Check if this content is safe according to our safety policy with the below categories.

#         # <BEGIN UNSAFE CONTENT CATEGORIES>
#         # {everyone_content_policy['policy']}
#         # <END UNSAFE CONTENT CATEGORIES>

#         # <CONTENT TO CHECK>
#         # {message}
#         # <END CONTENT TO CHECK>

#         # Provide your safety assessment for user in the above CONTENT TO CHECK section:
#         # - Reply ONLY with single word 'safe' or 'unsafe' [/INST]
#         # """

#         prompt = f"""[INST]Check if safe:
# {message}
# Reply 'safe' or 'unsafe'[/INST]
# """

#         # Generate safety check on cuda

#         # with torch.no_grad():
#         #     inputs = safety_tokenizer(
#         #         prompt,
#         #         return_tensors="pt",
#         #         padding=True,
#         #         truncation=True,
#         #     )

#         #     # Move inputs to correct device
#         #     inputs = {k: v.to(device) for k, v in inputs.items()}

#         #     output = safety_model.generate(
#         #         **inputs,
#         #         max_new_tokens=10,
#         #         temperature=0.1,  # Lower temperature for more consistent output
#         #         pad_token_id=safety_tokenizer.pad_token_id,  # Use configured pad token
#         #         eos_token_id=safety_tokenizer.eos_token_id,
#         #         do_sample=False,
#         #     )

#         # result = safety_tokenizer.decode(output[0], skip_special_tokens=True)
#         result = get_safety_response(prompt)
#         tprint(f"Raw safety check result: {result}")

#         # # Extract response after prompt
#         # if "[/INST]" in result:
#         #     result = result.split("[/INST]")[-1]

#         # # Clean response
#         # result = result.lower().strip()
#         # tprint(f"Cleaned safety check result: {result}")
#         # words = [word for word in result.split() if word in ["safe", "unsafe"]]

#         # # Take first valid response word
#         # is_safe = words[0] == "safe" if words else False

#         # tprint("Final Safety check result:", is_safe)

#         is_safe = "safe" in result.lower().split()

#         logger.info(
#             f"Safety check completed - Result: {'SAFE' if is_safe else 'UNSAFE'}"
#         )
#         return is_safe

#     except Exception as e:
#         logger.error(f"Safety check failed: {e}")
#         return False


# def detect_inventory_changes(game_state, output):
#     inventory = game_state["inventory"]
#     messages = [
#         {"role": "system", "content": system_prompt},
#         {"role": "user", "content": f"Current Inventory: {str(inventory)}"},
#         {"role": "user", "content": f"Recent Story: {output}"},
#         {"role": "user", "content": "Inventory Updates"},
#     ]

#     input_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
#     model_output = generator(input_text, num_return_sequences=1, temperature=0.0)
#     response = model_output[0]["generated_text"]
#     result = json.loads(response)
#     return result["itemUpdates"]


# def update_inventory(inventory, item_updates):
#     update_msg = ""
#     for update in item_updates:
#         name = update["name"]
#         change_amount = update["change_amount"]
#         if change_amount > 0:
#             if name not in inventory:
#                 inventory[name] = change_amount
#             else:
#                 inventory[name] += change_amount
#             update_msg += f"\nInventory: {name} +{change_amount}"
#         elif name in inventory and change_amount < 0:
#             inventory[name] += change_amount
#             update_msg += f"\nInventory: {name} {change_amount}"
#         if name in inventory and inventory[name] < 0:
#             del inventory[name]
#     return update_msg


logging.info("Finished helper function")