#4 SCENE GENERATION
from datetime import datetime
import os
import glob
import argparse
from typing import List
from utils import get_llm_response, read_yaml_to_string, extract_yaml_from_string
from prompt import SCENE_TEMPLATE


def generate_scene(worldview: str, character_profiles: List[str], model_name: str, temperature: float = 1.0) -> str:
    """Generate a scene based on worldview and character profiles.

    Args:
        worldview (str): World view description
        character_profiles (List[str]): List of character profile strings
        model_name (str): Name of the language model to use
        temperature (float, optional): Temperature parameter for generation. Defaults to 1.0.

    Returns:
        str: Generated scene in YAML format
    """
    # Combine all character information into a string
    combined_profiles = "\n\n".join(character_profiles)
    
    prompt = f"""
    The worldview of the scene is:
    {worldview}
    The characters in the scene include:
    {combined_profiles}
    There are no other characters in the scene except the above characters
    Please generate a scene based on the above information:
    Please strictly follow the following yaml format:
    ```yaml
    {SCENE_TEMPLATE}
    ```
    """
    response = get_llm_response(prompt, model_name, temperature=temperature)
    return extract_yaml_from_string(response)

def save_scene(scene: str, filename: str):
    """Save scene to a file.

    Args:
        scene (str): Scene content in YAML format
        filename (str): Path to the file to save the scene to

    Returns:
        None
    """
    with open(filename, "w", encoding="utf-8") as f:
        f.write(scene)

def load_character_profiles(paths: List[str]) -> List[str]:
    """Load multiple character profile files.

    Args:
        paths (List[str]): List of paths to character profile files, can include wildcards

    Returns:
        List[str]: List of character profile strings
    """
    profiles = []
    expanded_paths = []
    
    # Expand all wildcard paths
    for path in paths:
        path_expansions = glob.glob(path)
        if path_expansions:
            expanded_paths.extend(path_expansions)
        else:
            print(f"警告: 路径 '{path}' 未匹配到任何文件")
    
    # Ensure expanded_paths only contains YAML files
    yaml_files = [path for path in expanded_paths if path.endswith('.yaml') or path.endswith('.yml')]
    
    if not yaml_files:
        print(f"Warning: No YAML files found. The paths checked: {expanded_paths}")
        return profiles
        
    # Load each file
    for file_path in yaml_files:
        if os.path.isfile(file_path):
            try:
                profile = read_yaml_to_string(file_path)
                profiles.append(profile)
                print(f"Loaded character: {file_path}")
            except Exception as e:
                print(f"Failed to load character file {file_path}: {e}")
    
    return profiles

def get_default_character_paths():
    """Get default paths for character profile files.

    Args:
        None

    Returns:
        List[str]: List of default paths to character profile files
    """
    # Find the latest character_pool directory
    character_pools = glob.glob("character_pool/separated_*")
    if not character_pools:
        return ["character_pool/separated_*/Character*.yaml"]  # No directory found, return general path
    
    # Sort by timestamp, find the latest
    latest_pool = sorted(character_pools)[-1]
    return [f"{latest_pool}/*.yaml"]

if __name__ == "__main__":
    # Use 3 specified characters
    # python scene_generation.py --characters character_pool/separated_20250503_192001/Character1.yaml character_pool/separated_20250503_192001/Character2.yaml character_pool/separated_20250503_192001/Character3.yaml

    # # Use wildcard to select all characters
    # python scene_generation.py --characters "character_pool/separated_20250503_192001/*.yaml"

    # # Use different worldview and model
    # python scene_generation.py --worldview worldview/another_worldview.yaml --characters "character_pool/separated_20250503_192001/Character1.yaml" --model "gpt-4o" --temperature 0.8
    
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Generate a scene with specified characters')
    parser.add_argument('--worldview', type=str, default="worldview/worldview_20250503_185119.yaml",
                        help='Worldview file path')
    parser.add_argument('--characters', nargs='+', 
                        help='Character file path, can specify multiple file paths or use wildcards')
    parser.add_argument('--model', type=str, default="claude-3-7-sonnet-20250219",
                        help='Used model name')
    parser.add_argument('--temperature', type=float, default=1.0,
                        help='Temperature coefficient when generating')
    
    args = parser.parse_args()
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    try:
        # If no character file is specified, use default path
        if not args.characters:
            args.characters = get_default_character_paths()
            print(f"No character file specified, using default path: {args.characters}")
        
        # Read worldview file
        worldview_example = read_yaml_to_string(args.worldview)
        print(f"Loaded worldview file: {args.worldview}")
        
        # Read specified character files
        character_profiles = load_character_profiles(args.characters)
        
        if not character_profiles:
            print("No valid character files found, cannot generate scene")
            exit(1)
            
        print(f"Loaded {len(character_profiles)} characters")
        
        # Generate scene
        scene = generate_scene(worldview_example, character_profiles, args.model, temperature=args.temperature)
        print("\nGenerated scene:")
        print(scene[:500] + "..." if len(scene) > 500 else scene)

        # Save scene
        os.makedirs("scene", exist_ok=True)
        scene_filename = f"scene/scene_{timestamp}.yaml"
        save_scene(scene, scene_filename)

        print(f"Scene saved to {scene_filename}")
        
        # Record used character files
        record_file = f"scene/scene_{timestamp}_character_files.txt"
        with open(record_file, "w", encoding="utf-8") as f:
            f.write("\n".join(args.characters))
        print(f"Recorded used character files to {record_file}")
            
    except Exception as e:
        import traceback
        print(f"Error when generating scene: {e}")
        traceback.print_exc()