import os
import fnmatch
import re
from io import StringIO
import threading
import queue
from pygltflib import GLTF2
import librosa
import soundfile as sf
from PIL import Image
import random
import torch
import numpy as np

# Import for SamplingParams
try:
    from vllm import SamplingParams
except ImportError:
    # Define a simple SamplingParams class for use with the server
    class SamplingParams:
        def __init__(self, **kwargs):
            self.__dict__.update(kwargs)

def parse_gitignore(gitignore_path):
    """
    Parse a .gitignore file and return a list of patterns.
    
    Args:
        gitignore_path: Path to the .gitignore file
        
    Returns:
        List of gitignore patterns
    """
    patterns = []
    try:
        with open(gitignore_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                # Skip empty lines and comments
                if not line or line.startswith("#"):
                    continue
                patterns.append(line)
    except Exception as e:
        print(f"Warning: Error reading .gitignore file: {e}")
    
    return patterns

def is_ignored_by_gitignore(file_path, root_dir, gitignore_patterns):
    """
    Check if a file should be ignored based on gitignore patterns.
    
    Args:
        file_path: Absolute path to the file
        root_dir: Root directory of the project
        gitignore_patterns: List of patterns from .gitignore
        
    Returns:
        True if the file should be ignored, False otherwise
    """
    if not gitignore_patterns:
        return False
    
    # Get the relative path from the root directory
    rel_path = os.path.relpath(file_path, root_dir)
    # Use forward slashes for consistency (gitignore standard)
    rel_path = rel_path.replace(os.path.sep, "/")
    
    # Get the basename for matching
    basename = os.path.basename(file_path)
    
    # Check each pattern
    for pattern in gitignore_patterns:
        # Skip empty patterns
        if not pattern:
            continue
        
        # Handle negated patterns
        negated = pattern.startswith("!")
        if negated:
            pattern = pattern[1:]
        
        # Skip patterns that start with # (comments)
        if pattern.startswith("#"):
            continue
        
        # Handle directory-only pattern (trailing slash)
        dir_only = pattern.endswith("/")
        if dir_only:
            pattern = pattern[:-1]
            # Skip if file_path is not a directory
            if not os.path.isdir(file_path):
                continue
        
        # Simple exact match
        if rel_path == pattern or rel_path.startswith(pattern + "/"):
            return not negated
        
        # Handle glob patterns
        if "*" in pattern or "?" in pattern or "[" in pattern:
            # Try to match the pattern
            if fnmatch.fnmatch(rel_path, pattern):
                return not negated
            
            # Also check if the pattern matches just the basename
            if fnmatch.fnmatch(basename, pattern):
                return not negated
    
    return False

def process_asset_file(path, file_ext):
    """
    Process an asset file and return information about it based on its type.
    
    Args:
        path: Path to the file
        file_ext: File extension (lowercase)
        
    Returns:
        A string with information about the file
    """
    # Process image files
    if file_ext in ["png", "jpg", "jpeg"]:
        try:
            with Image.open(path) as img:
                width, height = img.size
                return f"({width}x{height}px)"
        except Exception as e:
            return f"(error: {str(e)})"

    # Process audio files
    elif file_ext in ["wav", "mp3", "ogg"]:
        audio, sr = librosa.load(path)
        audio_info = f"{librosa.beat.tempo(y=audio, sr=sr)[0]:.2f} BPM, {sf.info(path).duration:.2f}s"#, {sr}Hz"
        return f"({audio_info})"
    
    # Process 3D model files
    elif file_ext in ["glb"]:
        # Full mesh is loaded with default material automatically, we just need the animations if they exists
        gltf = GLTF2().load(path)
        animation_names = [animation.name for i, animation in enumerate(gltf.animations)]
        if len(animation_names) > 0:
            info_3d = f"animations: {', '.join(animation_names)}"
            return f"({info_3d})"
        else:
            return f""
    
    # Other allowed files
    return ""

def write_directory_tree(file, root_dir, exclude_dirs, prefix="", gitignore_patterns=None, project_root=None, allowed_asset_types=None, assets_structure="full"):
    """
    Write a directory tree structure to a file, excluding specified directories
    and respecting gitignore patterns.
    
    Args:
        file: File-like object to write to
        root_dir: Root directory to start from
        exclude_dirs: Set of directory names to exclude
        prefix: Prefix for indentation
        gitignore_patterns: List of patterns from .gitignore
        project_root: Root directory of the project (for gitignore pattern matching)
        allowed_asset_types: List of allowed asset file extensions
        assets_structure: How to display asset directory structure ('full' or 'compressed')
    """
    # Use default allowed asset types if none provided
    if allowed_asset_types is None:
        allowed_asset_types = ["png", "jpeg", "jpg", "ogg", "mp3", "wav", "glb"]
    # Use the root_dir as project_root if not provided
    if project_root is None:
        project_root = root_dir
    
    try:
        # Get items in the directory, excluding those in exclude_dirs
        items = []
        for item in sorted(os.listdir(root_dir)):
            if item in exclude_dirs:
                continue
            
            path = os.path.join(root_dir, item)
            
            # Skip items ignored by gitignore
            if gitignore_patterns and is_ignored_by_gitignore(path, project_root, gitignore_patterns):
                continue
            
            is_dir = os.path.isdir(path)
            items.append((item, path, is_dir))
    except PermissionError:
        file.write(f"{prefix}[Permission Denied]\n")
        return
    except FileNotFoundError:
        file.write(f"{prefix}[Directory Not Found]\n")
        return
    
    # For compressed mode, we need to check if this directory contains any files with allowed extensions
    if assets_structure == "compressed":
        has_allowed_files = False
        for _, path, is_dir in items:
            if not is_dir:
                item_lower = os.path.basename(path).lower()
                file_ext = item_lower.split('.')[-1] if '.' in item_lower else ""
                if file_ext in allowed_asset_types:
                    has_allowed_files = True
                    break
            else:
                # Recursively check subdirectories
                for root, _, files in os.walk(path):
                    for file_name in files:
                        file_ext = file_name.lower().split('.')[-1] if '.' in file_name else ""
                        if file_ext in allowed_asset_types:
                            has_allowed_files = True
                            break
                    if has_allowed_files:
                        break
            if has_allowed_files:
                break
        
        # If no allowed files in this directory or its subdirectories, don't show it
        if not has_allowed_files and items:
            return
    
    # Process each item
    for i, (item, path, is_dir) in enumerate(items):
        is_last = i == len(items) - 1
        
        # Write the item with appropriate connectors
        connector = "└── " if is_last else "├── "
        
        if is_dir:
            file.write(f"{prefix}{connector} {item}/\n")
            
            # If it's a directory, recursively write its contents
            next_prefix = prefix + ("    " if is_last else "│   ")
            
            if assets_structure == "compressed":
                # For compressed mode, check if this directory contains files directly
                has_direct_files = False
                for file_name in os.listdir(path):
                    file_path = os.path.join(path, file_name)
                    if not os.path.isdir(file_path):
                        file_ext = file_name.lower().split('.')[-1] if '.' in file_name else ""
                        if file_ext in allowed_asset_types:
                            has_direct_files = True
                            break
                
                if not has_direct_files:
                    # If no direct files, find the first subdirectory with files
                    subdirs = []
                    for subdir in sorted(os.listdir(path)):
                        subdir_path = os.path.join(path, subdir)
                        if os.path.isdir(subdir_path) and subdir not in exclude_dirs:
                            subdirs.append((subdir, subdir_path))
                    
                    if len(subdirs) == 1:
                        # If only one subdirectory, compress the tree
                        subdir, subdir_path = subdirs[0]
                        write_directory_tree(file, subdir_path, exclude_dirs, next_prefix, gitignore_patterns, project_root, allowed_asset_types, assets_structure)
                        continue
            
            # Normal recursive call for full mode or when compression isn't applicable
            write_directory_tree(file, path, exclude_dirs, next_prefix, gitignore_patterns, project_root, allowed_asset_types, assets_structure)
        else:
            # Process files based on their type
            item_lower = item.lower()
            file_ext = item_lower.split('.')[-1] if '.' in item_lower else ""
            
            # Only include files with allowed extensions
            if file_ext in allowed_asset_types:
                # Get file info using the common processing function
                file_info = process_asset_file(path, file_ext)
                file.write(f"{prefix}{connector} {item} {file_info}\n")

def scan_asset_directory(asset_dir, allowed_asset_types=None, respect_gitignore=True, assets_structure="full"):
    """
    Scan the asset directory and build a tree structure of assets with dimensions.
    
    Args:
        asset_dir: Path to the asset directory
        allowed_asset_types: List of allowed asset file extensions
        respect_gitignore: Whether to respect .gitignore patterns
        assets_structure: How to display asset directory structure ('full' or 'compressed')
        
    Returns:
        A string representation of the asset directory tree with file dimensions
    """
    if not asset_dir or not os.path.isdir(asset_dir):
        print("<< NO ASSET USED >>")
        return None
    
    # Use default allowed asset types if none provided
    if allowed_asset_types is None:
        allowed_asset_types = ["png", "jpeg", "jpg", "ogg", "mp3", "wav", "glb"]
    
    # Check for .gitignore file and parse patterns
    gitignore_patterns = []
    if respect_gitignore:
        gitignore_path = os.path.join(asset_dir, ".gitignore")
        if os.path.isfile(gitignore_path):
            gitignore_patterns = parse_gitignore(gitignore_path)
            print(f"Found .gitignore with {len(gitignore_patterns)} patterns - will respect these patterns")
    
    # Use StringIO to build the tree
    tree_file = StringIO()
    
    # Write the root directory
    tree_file.write(f"Assets Directory: {os.path.basename(asset_dir)}\n")
    tree_file.write("└── " + os.path.basename(asset_dir) + "/\n")
    
    # Define a set of directories to exclude
    exclude_dirs = {'.git', '__pycache__', 'node_modules', 'venv', '.env', 'env', 'dist', 'build'}
    
    # Write the directory tree
    write_directory_tree(tree_file, asset_dir, exclude_dirs, "    ", gitignore_patterns, asset_dir, allowed_asset_types, assets_structure)
    
    return tree_file.getvalue()

def select_assets_from_chunk(system_prompt, llm, content_description, pack_name, chunk_assets, max_assets_from_chunk, max_tokens, max_tries, seed, asset_dir, selection_idx, total_selections, temperature=0.7, top_p=0.95, top_k=-1, repetition_penalty=1.0):
    """
    Select assets from a chunk of assets within a pack.
    
    Args:
        system_prompt: System prompt for the LLM
        llm: The coding LLM
        content_description: Description of the content
        pack_name: Name of the asset pack
        chunk_assets: List of assets in this chunk
        max_assets_from_chunk: Maximum number of assets to select from this chunk
        max_tokens: Maximum number of tokens for generation
        max_tries: Maximum number of retry attempts
        seed: Random seed for generation
        asset_dir: Path to the asset directory
        selection_idx: Index of this selection (1-based)
        total_selections: Total number of selections being made
    
    Returns:
        List of selected asset paths from this chunk
    """
    print(f"Selecting from chunk {selection_idx}/{total_selections} with {len(chunk_assets)} assets")
    
    # Create a prompt for the model to choose specific assets from this chunk
    assets_prompt = f"""
You are helping to create a {content_description}.

This is selection {selection_idx} of {total_selections} from the sample pack "{pack_name}".
You need to choose at most {max_assets_from_chunk} assets from this portion of the pack that would be most useful for this content.
Here are the available assets from this portion:

"""
    
    # Group assets by type and organize into a tree structure
    assets_by_type = {}
    for i, asset in enumerate(chunk_assets):
        if asset["type"] not in assets_by_type:
            assets_by_type[asset["type"]] = []
        # Store the index for later reference
        asset_with_index = asset.copy()
        asset_with_index["index"] = i
        assets_by_type[asset["type"]].append(asset_with_index)
    
    # Build a tree structure for each asset type
    for i, (asset_type, assets) in enumerate(assets_by_type.items()):
        is_last_type = i == len(assets_by_type) - 1
        type_connector = "└── " if is_last_type else "├── "
        assets_prompt += f"\n{type_connector}{asset_type.upper()} ASSETS ({len(assets)} total):\n"
        
        # Group assets by their folder structure
        assets_tree = {}
        for asset in assets:
            path_parts = asset["path"].split(os.sep)
            current = assets_tree
            
            # Build the tree structure
            for part in path_parts[:-1]:  # Process directories
                if part not in current:
                    current[part] = {}
                current = current[part]
            
            # Add the file with its index
            filename = path_parts[-1]
            if "__files__" not in current:
                current["__files__"] = []
            current["__files__"].append((filename, asset["index"]))
        
        # Create a temporary list for building the tree
        tree_parts = []
        
        # Create a counter for sequential numbering
        sequential_counter = [0]
        
        # Create a mapping from sequential index to original asset index
        sequential_to_original_index = {}
        
        # Write the tree for this asset type
        type_prefix = "    " if is_last_type else "│   "
        
        # Define a function to recursively write the tree with proper connectors
        def write_asset_tree_to_list(tree, prefix="    ", path_prefix=""):
            items = sorted([(k, v) for k, v in tree.items() if k != "__files__"])
            files = sorted(tree.get("__files__", []))
            
            # Process all subdirectories
            for i, (name, contents) in enumerate(items):
                is_last_item = (i == len(items) - 1) and not files
                connector = prefix + ("└── " if is_last_item else "├── ")
                tree_parts.append(f"{connector}{name}/\n")
                
                # Determine the new prefix for the next level
                new_prefix = prefix + ("    " if is_last_item else "│   ")
                new_path_prefix = path_prefix + name + os.sep
                
                # Recursively process the subdirectory
                write_asset_tree_to_list(contents, new_prefix, new_path_prefix)
            
            # Process all files in the current directory
            for i, (filename, index) in enumerate(files):
                is_last_file = i == len(files) - 1
                connector = prefix + ("└── " if is_last_file else "├── ")
                # Use sequential counter instead of original index
                sequential_index = sequential_counter[0]
                # Store mapping from sequential index to original index
                sequential_to_original_index[sequential_index] = index
                sequential_counter[0] += 1
                
                # Get the full path to the asset
                full_path = os.path.join(asset_dir, chunk_assets[index]["path"])
                
                # Get file info using the common processing function
                file_ext = filename.lower().split('.')[-1] if '.' in filename else ""
                file_info = process_asset_file(full_path, file_ext)
                
                tree_parts.append(f"{connector}{filename} {file_info} [{sequential_index}]\n")
        
        # Write the tree for this asset type to the list
        write_asset_tree_to_list(assets_tree, type_prefix)
        
        # Append the tree parts to the assets_prompt
        assets_prompt += "".join(tree_parts)
    
    assets_prompt += f"""
Based on the content description, choose at most {max_assets_from_chunk} assets from this portion that would be most useful.
Choose a mix of sounds, music, and graphics (2D or 3D) based on the needs of the content.
Return your answer as a comma-separated list of NUMBERS (the indices in square brackets) within \\boxed{{}}, for example: "\\boxed{{0, 3, 5, 7}}"
"""
    print(assets_prompt)
    
    # Generate the response to choose specific assets with retry logic
    chunk_chosen_assets = []
    attempts = 0
    
    while not chunk_chosen_assets and attempts < max_tries:
        attempts += 1
        
        if attempts > 1:
            print(f"Retry attempt {attempts}/{max_tries} for chunk asset selection...")
            # Adjust the prompt for retry attempts to encourage better selection
            assets_prompt += f"\n\nPrevious selection was invalid. Please choose valid asset indices from the list above."
        
        if system_prompt is not None:
            messages = [{"role": "system", "content": system_prompt}]
        else:
            messages = []
        
        messages.append({"role": "user", "content": assets_prompt})
        
        sampling_params = SamplingParams(
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
            max_tokens=max_tokens,
            seed=seed
        )
        
        # Increment seed for next generation
        seed += 1
        
        response = llm.generate(messages, sampling_params)
        chosen_indices_text = response[0].outputs[0].text.strip()
        print(f"Model response: {chosen_indices_text}")
        
        # Parse the chosen indices
        try:
            # Extract all boxed content if present and take the last one
            boxed_matches = re.findall(r'\\boxed\{(.*?)\}', chosen_indices_text)
            if boxed_matches:
                # Use the last boxed content
                last_boxed = boxed_matches[-1]
                # Extract numbers from the last boxed content
                indices = re.findall(r'\d+', last_boxed)
            else:
                # Fallback to extracting numbers from the entire response
                indices = re.findall(r'\d+', chosen_indices_text)
            
            chosen_indices = [int(idx) for idx in indices if int(idx) < len(chunk_assets)]
            
            # Limit to max_assets_from_chunk
            chosen_indices = chosen_indices[:min(max_assets_from_chunk, len(chosen_indices))]
            
            if chosen_indices:
                # Get the actual asset paths for the chosen indices using our mapping
                chunk_chosen_assets = []
                for idx in chosen_indices:
                    if idx in sequential_to_original_index:
                        original_idx = sequential_to_original_index[idx]
                        chunk_chosen_assets.append(chunk_assets[original_idx]["path"])
                print(f"Successfully selected {len(chunk_chosen_assets)} assets from chunk {selection_idx} on attempt {attempts}/{max_tries}")
                print(f"Selected assets: {chunk_chosen_assets}")
        except Exception as e:
            print(f"Error parsing indices: {e}")
            chunk_chosen_assets = []
    
    # If all attempts failed, fall back to the first few available assets
    if not chunk_chosen_assets:
        print(f"Warning: Model failed to choose valid assets after {max_tries} attempts, using the first few available assets from chunk {selection_idx}")
        chunk_chosen_assets = [chunk_assets[i]["path"] for i in range(min(max_assets_from_chunk, len(chunk_assets)))]
    
    return chunk_chosen_assets

def select_assets_with_model(system_prompt, llm, content_description, asset_dir, allowed_asset_types=None, max_sample_packs=2, max_assets=20, max_tokens=1024, max_tries=10, seed=42, select_sample_packs_only=False, enable_audio=False, assets_selection="individual", max_assets_per_pack=None, temperature=0.7, top_p=0.95, top_k=-1, repetition_penalty=1.0):
    """
    Let the coding model choose asset packs and specific assets based on the game description.
    
    Args:
        system_prompt: System prompt for the LLM
        llm: The coding LLM
        content_description: Description of the content
        asset_dir: Path to the asset directory
        allowed_asset_types: List of allowed asset file extensions
        max_sample_packs: Maximum number of asset packs to select
        max_assets: Maximum number of assets to select
        max_tries: Maximum number of retry attempts if asset selection fails
        seed: Random seed for generation
        select_sample_packs_only: If True, only select asset packs and use all assets in those packs
        
    Returns:
        A tuple containing:
        - A string representation of the selected asset tree with dimensions
        - A list of selected asset paths
        - Updated seed value
    """
    if not asset_dir or not os.path.isdir(asset_dir):
        print("<< NO ASSET USED >>")
        return None, [], seed
    
    # Use default allowed asset types if none provided
    if allowed_asset_types is None:
        allowed_asset_types = ["png", "jpeg", "jpg", "ogg", "mp3", "wav", "glb"]
    
    # Step 1: Get all asset packs (folders) in the asset directory with file counts by type
    asset_packs = []
    for folder in os.listdir(asset_dir):
        folder_path = os.path.join(asset_dir, folder)
        if os.path.isdir(folder_path) and not folder.startswith('.'):
            # Count files by type in this folder (recursively)
            file_counts = {}
            for root, _, files in os.walk(folder_path):
                for file in files:
                    ext = file.lower().split('.')[-1] if '.' in file else ""
                    if ext in allowed_asset_types:
                        file_counts[ext] = file_counts.get(ext, 0) + 1
            
            # Add to asset packs list if it contains any allowed files
            if file_counts:
                asset_packs.append({
                    "name": folder,
                    "file_counts": file_counts,
                    "total_files": sum(file_counts.values())
                })
    
    if not asset_packs:
        print(f"Warning: No asset packs found in {asset_dir}")
        return None, [], seed
    
    modalities_text = " It should integrate visual and auditory elements." if enable_audio else " It should integrate visual elements."
    audio_text = " Ensure that audio assets are included." if enable_audio else ""

    # Create a prompt for the model to choose asset packs
    asset_packs_prompt = f"""
You are helping to create the following content: {content_description}.{modalities_text}

You need to choose at most {max_sample_packs} asset packs (folders) that would be most useful for this content.
Here are the available asset packs with the number of files by type:

"""
    
    for i, pack in enumerate(asset_packs):
        is_last = i == len(asset_packs) - 1
        connector = "└── " if is_last else "├── "
        asset_packs_prompt += f"{connector}{pack['name']} ({pack['total_files']} total files):\n"
        
        for j, (ext, count) in enumerate(pack['file_counts'].items()):
            is_last_ext = j == len(pack['file_counts']) - 1
            ext_connector = "    └── " if is_last_ext else "    ├── "
            asset_packs_prompt += f"{ext_connector}{ext}: {count} files\n"
    
    asset_packs_prompt += f"""
Based on the content description, choose at most {max_sample_packs} asset packs that would be most useful.{audio_text}
Return your answer as a comma-separated list of folder names within \\boxed{{}}, for example: "\\boxed{{folder1, folder2, folder3}}"
"""
    print(asset_packs_prompt)
    
    # Generate the response to choose asset packs with retry logic
    chosen_packs = []
    attempts = 0
    
    while not chosen_packs and attempts < max_tries:
        attempts += 1
        
        if attempts > 1:
            print(f"Retry attempt {attempts}/{max_tries} for asset pack selection...")
            # Adjust the prompt for retry attempts to encourage better selection
            asset_packs_prompt += f"\n\nPrevious selection was invalid. Please choose valid folder names from the list above."
        
        if system_prompt is not None:
            messages = [{"role": "system", "content": system_prompt}]
        else:
            messages = []
        
        messages.append({"role": "user", "content": asset_packs_prompt})
        
        sampling_params = SamplingParams(
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
            max_tokens=max_tokens,
            seed=seed
        )
        
        # Increment seed for next generation
        seed += 1
        
        response = llm.generate(messages, sampling_params)
        chosen_packs_text = response[0].outputs[0].text.strip()
        
        # Parse the chosen packs
        # Extract all boxed content if present and take the last one
        boxed_matches = re.findall(r'\\boxed\{(.*?)\}', chosen_packs_text)
        if boxed_matches:
            # Use the last boxed content
            last_boxed = boxed_matches[-1]
            chosen_packs_text = last_boxed
        
        chosen_packs = [pack.strip() for pack in chosen_packs_text.split(',')]
        chosen_packs = [pack for pack in chosen_packs if pack in [p["name"] for p in asset_packs]]
        
        # Limit to max_sample_packs
        chosen_packs = chosen_packs[:min(max_sample_packs, len(chosen_packs))]
        
        if chosen_packs:
            print(f"Successfully selected asset packs on attempt {attempts}/{max_tries}")
            
            # Validate chosen packs against specific rules
            valid_selection = True
            rejection_reason = ""
            
            # Rule 1: If enable_audio=True, at least one pack should contain audio
            if enable_audio:
                has_audio = False
                for pack_name in chosen_packs:
                    pack_info = next((p for p in asset_packs if p["name"] == pack_name), None)
                    if pack_info and any(ext in pack_info["file_counts"] for ext in ["wav", "mp3", "ogg"]):
                        has_audio = True
                        break
                
                if not has_audio:
                    valid_selection = False
                    rejection_reason = "No audio assets found in selected packs, but audio is required"
            
            # Rule 2: If "3D" is in content_description, at least one pack should contain 3D models
            if "3D" in content_description:
                has_3d_models = False
                for pack_name in chosen_packs:
                    pack_info = next((p for p in asset_packs if p["name"] == pack_name), None)
                    if pack_info and "glb" in pack_info["file_counts"]:
                        has_3d_models = True
                        break
                
                if not has_3d_models:
                    valid_selection = False
                    rejection_reason = "No 3D models found in selected packs, but 3D content is required"
            
            # Rule 3: In all cases, at least one pack should contain images
            has_images = False
            for pack_name in chosen_packs:
                pack_info = next((p for p in asset_packs if p["name"] == pack_name), None)
                if pack_info and any(ext in pack_info["file_counts"] for ext in ["png", "jpg", "jpeg"]):
                    has_images = True
                    break
            
            if not has_images:
                valid_selection = False
                rejection_reason = "No image assets found in selected packs, but images are required"
            
            # If any rule is broken, reject the chosen packs and retry
            if not valid_selection:
                print(f"Rejecting chosen packs: {rejection_reason}")
                chosen_packs = []  # Reset chosen_packs to trigger another iteration
                asset_packs_prompt += f"\n\nPrevious selection was rejected: {rejection_reason}. Please choose different packs that meet all requirements."
    
    # If all attempts failed, fall back to the first available pack
    if not chosen_packs:
        print(f"Warning: Model failed to choose valid asset packs after {max_tries} attempts, using the first available pack")
        chosen_packs = [asset_packs[0]["name"]]
    
    print(f"Model selected {len(chosen_packs)}/{len(asset_packs)} asset packs: {', '.join(chosen_packs)}")
    
    # Step 2: Get all assets in the chosen packs
    all_assets = []
    for pack in chosen_packs:
        pack_path = os.path.join(asset_dir, pack)
        for root, _, files in os.walk(pack_path):
            for file in files:
                ext = file.lower().split('.')[-1] if '.' in file else ""
                if ext in allowed_asset_types:
                    # Get the relative path from the asset directory
                    rel_path = os.path.relpath(os.path.join(root, file), asset_dir)
                    file_type = "image" if ext in ["png", "jpg", "jpeg"] else "audio" if ext in ["wav", "mp3", "ogg"] else "3D model" if ext == "glb" else "other"
                    all_assets.append({
                        "path": rel_path,
                        "type": file_type,
                        "extension": ext
                    })
    
    if not all_assets:
        print(f"Warning: No assets found in the chosen packs")
        return None, [], seed
    
    # If select_sample_packs_only is True, use all assets from the chosen packs
    chosen_assets = []
    if select_sample_packs_only:
        print(f"Using all assets from the selected packs (--select_sample_packs mode)")
        chosen_assets = [asset["path"] for asset in all_assets]
        print(f"Using all {len(chosen_assets)} assets from the selected packs")
    else:
        # Group assets by their sample pack
        assets_by_pack = {}
        for asset in all_assets:
            pack_name = asset["path"].split(os.sep)[0]
            if pack_name not in assets_by_pack:
                assets_by_pack[pack_name] = []
            assets_by_pack[pack_name].append(asset)
        
        # Calculate max assets per pack to stay within overall max_assets limit
        max_per_pack = max(1, max_assets // len(assets_by_pack))
        print(f"Selecting at most {max_per_pack} assets per pack")
        
        if assets_selection == "combined":
            print(f"Using combined asset selection mode to select assets from all packs at once")
            
            # Create a prompt for the model to choose specific assets from all packs
            assets_prompt = f"""
You are helping to create a {content_description}.

You need to choose at most {max_assets} assets from all sample packs that would be most useful for this content.
Here are the available assets from all sample packs:

"""
            
            # Group assets by type and organize into a tree structure
            assets_by_type = {}
            for i, asset in enumerate(all_assets):
                if asset["type"] not in assets_by_type:
                    assets_by_type[asset["type"]] = []
                # Store the index for later reference
                asset_with_index = asset.copy()
                asset_with_index["index"] = i
                assets_by_type[asset["type"]].append(asset_with_index)
            
            # Build a tree structure for each asset type
            for i, (asset_type, assets) in enumerate(assets_by_type.items()):
                is_last_type = i == len(assets_by_type) - 1
                type_connector = "└── " if is_last_type else "├── "
                assets_prompt += f"\n{type_connector}{asset_type.upper()} ASSETS ({len(assets)} total):\n"
                
                # Group assets by their folder structure
                assets_tree = {}
                for asset in assets:
                    path_parts = asset["path"].split(os.sep)
                    current = assets_tree
                    
                    # Build the tree structure
                    for part in path_parts[:-1]:  # Process directories
                        if part not in current:
                            current[part] = {}
                        current = current[part]
                    
                    # Add the file with its index
                    filename = path_parts[-1]
                    if "__files__" not in current:
                        current["__files__"] = []
                    current["__files__"].append((filename, asset["index"]))
                
                # Create a temporary list for building the tree
                tree_parts = []
                
                # Create a counter for sequential numbering
                sequential_counter = [0]
                
                # Create a mapping from sequential index to original asset index
                sequential_to_original_index = {}
                
                # Define a function to recursively write the tree with proper connectors
                def write_asset_tree_to_list(tree, prefix="    ", path_prefix=""):
                    items = sorted([(k, v) for k, v in tree.items() if k != "__files__"])
                    files = sorted(tree.get("__files__", []))
                    
                    # Process all subdirectories
                    for i, (name, contents) in enumerate(items):
                        is_last_item = (i == len(items) - 1) and not files
                        connector = prefix + ("└── " if is_last_item else "├── ")
                        tree_parts.append(f"{connector}{name}/\n")
                        
                        # Determine the new prefix for the next level
                        new_prefix = prefix + ("    " if is_last_item else "│   ")
                        new_path_prefix = path_prefix + name + os.sep
                        
                        # Recursively process the subdirectory
                        write_asset_tree_to_list(contents, new_prefix, new_path_prefix)
                    
                    # Process all files in the current directory
                    for i, (filename, index) in enumerate(files):
                        is_last_file = i == len(files) - 1
                        connector = prefix + ("└── " if is_last_file else "├── ")
                        # Use sequential counter instead of original index
                        sequential_index = sequential_counter[0]
                        # Store mapping from sequential index to original index
                        sequential_to_original_index[sequential_index] = index
                        sequential_counter[0] += 1
                        
                        # Get the full path to the asset
                        full_path = os.path.join(asset_dir, all_assets[index]["path"])
                        
                        # Get file info using the common processing function
                        file_ext = filename.lower().split('.')[-1] if '.' in filename else ""
                        file_info = process_asset_file(full_path, file_ext)
                        
                        tree_parts.append(f"{connector}{filename} {file_info} [{sequential_index}]\n")
                
                # Write the tree for this asset type
                type_prefix = "    " if is_last_type else "│   "
                
                # Write the tree for this asset type to the list
                write_asset_tree_to_list(assets_tree, type_prefix)
                
                # Append the tree parts to the assets_prompt
                assets_prompt += "".join(tree_parts)
            
            assets_prompt += f"""
Based on the content description, choose at most {max_assets} assets from all sample packs that would be most useful.
Choose a mix of sounds, music, and graphics (2D or 3D) based on the needs of the content.
Return your answer as a comma-separated list of NUMBERS (the indices in square brackets) within \\boxed{{}}, for example: "\\boxed{{0, 3, 5, 7}}"
"""
            print(assets_prompt)
            
            # Generate the response to choose specific assets with retry logic
            combined_chosen_assets = []
            attempts = 0
            
            while not combined_chosen_assets and attempts < max_tries:
                attempts += 1
                
                if attempts > 1:
                    print(f"Retry attempt {attempts}/{max_tries} for asset selection...")
                    # Adjust the prompt for retry attempts to encourage better selection
                    assets_prompt += f"\n\nPrevious selection was invalid. Please choose valid asset indices from the list above."
                
                if system_prompt is not None:
                    messages = [{"role": "system", "content": system_prompt}]
                else:
                    messages = []
                
                messages.append({"role": "user", "content": assets_prompt})
                
                sampling_params = SamplingParams(
                    temperature=temperature,
                    top_p=top_p,
                    top_k=top_k,
                    repetition_penalty=repetition_penalty,
                    max_tokens=max_tokens,
                    seed=seed
                )
                
                # Increment seed for next generation
                seed += 1
                
                response = llm.generate(messages, sampling_params)
                chosen_indices_text = response[0].outputs[0].text.strip()
                print(f"Model response: {chosen_indices_text}")
                
                # Parse the chosen indices
                try:
                    # Extract all boxed content if present and take the last one
                    boxed_matches = re.findall(r'\\boxed\{(.*?)\}', chosen_indices_text)
                    if boxed_matches:
                        # Use the last boxed content
                        last_boxed = boxed_matches[-1]
                        # Extract numbers from the last boxed content
                        indices = re.findall(r'\d+', last_boxed)
                    else:
                        # Fallback to extracting numbers from the entire response
                        indices = re.findall(r'\d+', chosen_indices_text)
                    
                    chosen_indices = [int(idx) for idx in indices if int(idx) < len(all_assets)]
                    
                    # Limit to max_assets
                    chosen_indices = chosen_indices[:min(max_assets, len(chosen_indices))]
                    
                    if chosen_indices:
                        # Get the actual asset paths for the chosen indices using our mapping
                        combined_chosen_assets = []
                        for idx in chosen_indices:
                            if idx in sequential_to_original_index:
                                original_idx = sequential_to_original_index[idx]
                                combined_chosen_assets.append(all_assets[original_idx]["path"])
                        print(f"Successfully selected {len(combined_chosen_assets)} assets across all packs on attempt {attempts}/{max_tries}")
                        print(f"Selected assets: {combined_chosen_assets}")
                except Exception as e:
                    print(f"Error parsing indices: {e}")
                    combined_chosen_assets = []
            
            # If all attempts failed, fall back to the first few available assets
            if not combined_chosen_assets:
                print(f"Warning: Model failed to choose valid assets after {max_tries} attempts, using the first few available assets")
                combined_chosen_assets = [all_assets[i]["path"] for i in range(min(max_assets, len(all_assets)))]
            
            # Add the chosen assets to the overall list
            chosen_assets = combined_chosen_assets
            
        else:
            # Process each sample pack individually (original behavior)
            print(f"Processing each sample pack individually to select assets")
            
            # Process each pack individually
            for pack_name, pack_assets in assets_by_pack.items():
                print(f"\nProcessing sample pack: {pack_name} ({len(pack_assets)} assets)")
                
                # Check if we need to split the pack due to max_assets_per_pack limit
                if max_assets_per_pack and len(pack_assets) > max_assets_per_pack:
                    print(f"Pack has {len(pack_assets)} assets which exceeds max_assets_per_pack ({max_assets_per_pack})")
                    print(f"Splitting pack into multiple selections...")
                    
                    # Calculate how many selections we need to split equally
                    num_selections = (len(pack_assets) + max_assets_per_pack - 1) // max_assets_per_pack
                    # Calculate chunk size to split as equally as possible
                    chunk_size = (len(pack_assets) + num_selections - 1) // num_selections
                    assets_per_selection = max_per_pack // num_selections
                    
                    print(f"Will make {num_selections} selections of ~{chunk_size} assets each (max {assets_per_selection} selected from each)")
                    
                    # Split assets into chunks of equal size
                    pack_chosen_assets = []
                    for selection_idx in range(num_selections):
                        start_idx = selection_idx * chunk_size
                        end_idx = min(start_idx + chunk_size, len(pack_assets))
                        chunk_assets = pack_assets[start_idx:end_idx]
                        
                        print(f"\nSelection {selection_idx + 1}/{num_selections}: Processing assets {start_idx + 1}-{end_idx} from pack {pack_name}")
                        
                        # If this chunk has few enough assets, take them all
                        if len(chunk_assets) <= assets_per_selection:
                            print(f"Chunk has {len(chunk_assets)} assets which is <= assets_per_selection ({assets_per_selection}), using all assets from this chunk")
                            chunk_chosen_assets = [asset["path"] for asset in chunk_assets]
                            pack_chosen_assets.extend(chunk_chosen_assets)
                        else:
                            # Let the model select from this chunk
                            chunk_chosen_assets = select_assets_from_chunk(
                                system_prompt, llm, content_description, pack_name, chunk_assets, 
                                assets_per_selection, max_tokens, max_tries, seed, asset_dir,
                                selection_idx + 1, num_selections,
                                temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty
                            )
                            pack_chosen_assets.extend(chunk_chosen_assets)
                            seed += max_tries  # Increment seed to account for potential retries
                    
                    chosen_assets.extend(pack_chosen_assets)
                    continue
                
                # If there are at least max_per_pack assets in the pack, automatically take all assets
                if len(pack_assets) <= max_per_pack:
                    print(f"Pack has {len(pack_assets)} assets which is <= max_per_pack ({max_per_pack}), automatically using all assets")
                    pack_chosen_assets = [asset["path"] for asset in pack_assets]
                    chosen_assets.extend(pack_chosen_assets)
                    continue
            
                # Create a prompt for the model to choose specific assets from this pack
                assets_prompt = f"""
    You are helping to create a {content_description}.

    You need to choose at most {max_per_pack} assets from the sample pack "{pack_name}" that would be most useful for this content.
    Here are the available assets from this sample pack:

    """
                
                # Group assets by type and organize into a tree structure
                assets_by_type = {}
                for i, asset in enumerate(pack_assets):
                    if asset["type"] not in assets_by_type:
                        assets_by_type[asset["type"]] = []
                    # Store the index for later reference
                    asset_with_index = asset.copy()
                    asset_with_index["index"] = i
                    assets_by_type[asset["type"]].append(asset_with_index)
                
                # Build a tree structure for each asset type
                for i, (asset_type, assets) in enumerate(assets_by_type.items()):
                    is_last_type = i == len(assets_by_type) - 1
                    type_connector = "└── " if is_last_type else "├── "
                    assets_prompt += f"\n{type_connector}{asset_type.upper()} ASSETS ({len(assets)} total):\n"
                    
                    # Group assets by their folder structure
                    assets_tree = {}
                    for asset in assets:
                        path_parts = asset["path"].split(os.sep)
                        current = assets_tree
                        
                        # Build the tree structure
                        for part in path_parts[:-1]:  # Process directories
                            if part not in current:
                                current[part] = {}
                            current = current[part]
                        
                        # Add the file with its index
                        filename = path_parts[-1]
                        if "__files__" not in current:
                            current["__files__"] = []
                        current["__files__"].append((filename, asset["index"]))
                    
                    # Create a temporary list for building the tree
                    tree_parts = []
                    
                    # Create a counter for sequential numbering
                    sequential_counter = [0]
                    
                    # Create a mapping from sequential index to original asset index
                    sequential_to_original_index = {}
                    
                    # Write the tree for this asset type
                    type_prefix = "    " if is_last_type else "│   "
                    
                    # Define a modified write_asset_tree function that appends to tree_parts
                    def write_asset_tree_to_list(tree, prefix="    ", path_prefix=""):
                        items = sorted([(k, v) for k, v in tree.items() if k != "__files__"])
                        files = sorted(tree.get("__files__", []))
                        
                        # Process all subdirectories
                        for i, (name, contents) in enumerate(items):
                            is_last_item = (i == len(items) - 1) and not files
                            connector = prefix + ("└── " if is_last_item else "├── ")
                            tree_parts.append(f"{connector}{name}/\n")
                            
                            # Determine the new prefix for the next level
                            new_prefix = prefix + ("    " if is_last_item else "│   ")
                            new_path_prefix = path_prefix + name + os.sep
                            
                            # Recursively process the subdirectory
                            write_asset_tree_to_list(contents, new_prefix, new_path_prefix)
                        
                        # Process all files in the current directory
                        for i, (filename, index) in enumerate(files):
                            is_last_file = i == len(files) - 1
                            connector = prefix + ("└── " if is_last_file else "├── ")
                            # Use sequential counter instead of original index
                            sequential_index = sequential_counter[0]
                            # Store mapping from sequential index to original index
                            sequential_to_original_index[sequential_index] = index
                            sequential_counter[0] += 1
                            
                            # Get the full path to the asset
                            full_path = os.path.join(asset_dir, pack_assets[index]["path"])
                            
                            # Get file info using the common processing function
                            file_ext = filename.lower().split('.')[-1] if '.' in filename else ""
                            file_info = process_asset_file(full_path, file_ext)
                            
                            tree_parts.append(f"{connector}{filename} {file_info} [{sequential_index}]\n")
                    
                    # Write the tree for this asset type to the list
                    write_asset_tree_to_list(assets_tree, type_prefix)
                    
                    # Append the tree parts to the assets_prompt
                    assets_prompt += "".join(tree_parts)
                
                assets_prompt += f"""
    Based on the content description, choose at most {max_per_pack} assets from this sample pack that would be most useful.
    Choose a mix of sounds, music, and graphics (2D or 3D) based on the needs of the content.
    Return your answer as a comma-separated list of NUMBERS (the indices in square brackets) within \\boxed{{}}, for example: "\\boxed{{0, 3, 5, 7}}"
    """
                print(assets_prompt)
                
                # Generate the response to choose specific assets with retry logic
                pack_chosen_assets = []
                attempts = 0
                
                while not pack_chosen_assets and attempts < max_tries:
                    attempts += 1
                    
                    if attempts > 1:
                        print(f"Retry attempt {attempts}/{max_tries} for asset selection...")
                        # Adjust the prompt for retry attempts to encourage better selection
                        assets_prompt += f"\n\nPrevious selection was invalid. Please choose valid asset indices from the list above."
                    
                    if system_prompt is not None:
                        messages = [{"role": "system", "content": system_prompt}]
                    else:
                        messages = []
                    
                    messages.append({"role": "user", "content": assets_prompt})
                    
                    sampling_params = SamplingParams(
                        temperature=temperature,
                        top_p=top_p,
                        top_k=top_k,
                        repetition_penalty=repetition_penalty,
                        max_tokens=max_tokens,
                        seed=seed
                    )
                    
                    # Increment seed for next generation
                    seed += 1
                    
                    response = llm.generate(messages, sampling_params)
                    chosen_indices_text = response[0].outputs[0].text.strip()
                    print(f"Model response: {chosen_indices_text}")
                    
                    # Parse the chosen indices
                    try:
                        # Extract all boxed content if present and take the last one
                        boxed_matches = re.findall(r'\\boxed\{(.*?)\}', chosen_indices_text)
                        if boxed_matches:
                            # Use the last boxed content
                            last_boxed = boxed_matches[-1]
                            # Extract numbers from the last boxed content
                            indices = re.findall(r'\d+', last_boxed)
                        else:
                            # Fallback to extracting numbers from the entire response
                            indices = re.findall(r'\d+', chosen_indices_text)
                        
                        chosen_indices = [int(idx) for idx in indices if int(idx) < len(pack_assets)]
                        
                        # Limit to max_per_pack
                        chosen_indices = chosen_indices[:min(max_per_pack, len(chosen_indices))]
                        
                        if chosen_indices:
                            # Get the actual asset paths for the chosen indices using our mapping
                            pack_chosen_assets = []
                            for idx in chosen_indices:
                                if idx in sequential_to_original_index:
                                    original_idx = sequential_to_original_index[idx]
                                    pack_chosen_assets.append(pack_assets[original_idx]["path"])
                            print(f"Successfully selected {len(pack_chosen_assets)} assets from pack {pack_name} on attempt {attempts}/{max_tries}")
                            print(f"Selected assets: {pack_chosen_assets}")
                    except Exception as e:
                        print(f"Error parsing indices: {e}")
                        pack_chosen_assets = []
                
                # If all attempts failed, fall back to the first few available assets
                if not pack_chosen_assets:
                    print(f"Warning: Model failed to choose valid assets after {max_tries} attempts, using the first few available assets from pack {pack_name}")
                    pack_chosen_assets = [pack_assets[i]["path"] for i in range(min(max_per_pack, len(pack_assets)))]
                
                # Add the chosen assets from this pack to the overall list
                chosen_assets.extend(pack_chosen_assets)
        
        print(f"Model selected {len(chosen_assets)}/{len(all_assets)} assets across all packs")
    
    # Step 3: Build the tree of selected assets
    tree_file = StringIO()
    tree_file.write(f"Directory structure:\n")
    tree_file.write("└── " + os.path.basename(asset_dir) + "/\n")
    
    # Build a nested dictionary structure representing the directory tree
    tree_structure = {}
    for asset_path in chosen_assets:
        parts = asset_path.split(os.sep)
        current = tree_structure
        for i, part in enumerate(parts[:-1]):  # Process directories
            if part not in current:
                current[part] = {}
            current = current[part]
        
        # Add the file (last part) with its full path for later processing
        if "__files__" not in current:
            current["__files__"] = []
        current["__files__"].append(asset_path)
    
    # Recursive function to write the tree with proper connectors
    def write_tree(file, structure, prefix="    ", is_last=False):
        items = sorted([(k, v) for k, v in structure.items() if k != "__files__"])
        files = structure.get("__files__", [])
        
        # Process all subdirectories
        for i, (name, contents) in enumerate(items):
            is_last_item = (i == len(items) - 1) and not files
            connector = "└── " if is_last_item else "├── "
            file.write(f"{prefix}{connector}{name}/\n")
            
            # Determine the new prefix for the next level
            new_prefix = prefix + ("    " if is_last_item else "│   ")
            
            # Recursively process the subdirectory
            write_tree(file, contents, new_prefix, is_last_item)
        
        # Process all files in the current directory
        for i, asset_path in enumerate(sorted(files)):
            is_last_file = i == len(files) - 1
            connector = "└── " if is_last_file else "├── "
            
            # Get just the filename (last part of the path)
            filename = os.path.basename(asset_path)
            
            # Get the full path to the asset
            full_path = os.path.join(asset_dir, asset_path)
            
            # Process the asset based on its type
            file_ext = filename.lower().split('.')[-1] if '.' in filename else ""
            
            # Get file info using the common processing function
            file_info = process_asset_file(full_path, file_ext)
            file.write(f"{prefix}{connector}{filename} {file_info}\n")
    
    # Write the tree structure starting with the top-level folders
    write_tree(tree_file, tree_structure)
    
    return tree_file.getvalue(), chosen_assets, seed
