#5 INITIAL DIALOGUE 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, read_json_to_string, extract_json_from_string
import json

DIALOGUE_TEMPLATE = read_json_to_string("prompt_template/dialogue.json")

def generate_dialogue(scene: str, worldview: str, character_profiles: List[str], model_name: str, temperature: float = 1.0, character_names: List[str] = None) -> str:
    """Generate a dialogue between characters in a scene.

    Args:
        scene (str): Scene description
        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.
        character_names (List[str], optional): List of character names. Defaults to None.

    Returns:
        str: Generated dialogue in JSON format
    """
    # Combine all character information into a string
    combined_profiles = "\n\n".join(character_profiles)
    
    # If no character name is provided, extract from character_profiles
    if not character_names:
        character_names = []
        for profile in character_profiles:
            for line in profile.split('\n'):
                if line.strip().startswith('Name:'):
                    name_value = line.strip()[5:].strip()
                    if name_value.startswith('"') and name_value.endswith('"'):
                        character_names.append(name_value[1:-1])
                    else:
                        character_names.append(name_value)
                    break
    
    character_names_str = ", ".join([f'"{name}"' for name in character_names if name])
    
    prompt = f"""
    The worldview of the scene is:
    {worldview}
    The characters in the scene include:
    {combined_profiles}
    The specific scene is:
    {scene}
    Please generate a dialogue between characters in the scene based on the above information, the characters' behavior should be consistent with their character_profile and consider the worldview and current scene information.
    For each character, [psychological activity] represents the character's inner thoughts, (action) represents the actual actions taken by the character, and the spoken content does not need to be symbolized.
    The content after each character_name can only be one of the three, or any combination of number and order.
    Additionally, note that the dialogue needs to insert environment, environment as a special character, used to objectively describe the real-time changes that may occur in the environment of the scene and the impact of the character's (action) interacting with items in the scene.
    
    The characters in the scene are: {character_names_str}
    
    Please strictly follow the following format, note that the JSON format is correct, each key-value pair must have a comma, and the last key-value pair must not have a comma:
    ```json
    {DIALOGUE_TEMPLATE}
    ```
    Special note: the character_name key needs to be replaced with the Name in the actual character_profile.
    """
    
    # Maximum of 3 attempts to generate valid JSON
    max_attempts = 3
    for attempt in range(max_attempts):
        try:
            response = get_llm_response(prompt, model_name, temperature=temperature)
            result = extract_json_from_string(response)
            
            # 检查结果类型
            if isinstance(result, str):
                print(f"Warning: Returned a string instead of a JSON object, try to manually parse")
                try:
                    return json.dumps(json.loads(result), ensure_ascii=False, indent=4)
                except:
                    return result
            elif isinstance(result, (dict, list)):
                return json.dumps(result, ensure_ascii=False, indent=4)
            else:
                return json.dumps({"error": "Unable to parse the result type", "raw": str(result)[:1000]}, ensure_ascii=False, indent=4)
                
        except Exception as e:
            print(f"Attempt {attempt+1}/{max_attempts} failed: {e}")
            if attempt == max_attempts - 1:
                # Last attempt failed, return error information
                return json.dumps({"error": f"无法生成有效对话: {str(e)}", "attempts": attempt+1}, ensure_ascii=False, indent=4)
    
    # Should not reach here
    return json.dumps({"error": "未知错误"}, ensure_ascii=False, indent=4)

def save_dialogue(dialogue: str, filename: str):
    """Save dialogue to a file.

    Args:
        dialogue (str): Dialogue content in JSON format or as a string
        filename (str): Path to the file to save the dialogue to

    Returns:
        bool: True if successful, False otherwise
    """
    try:
        # If the input is a dictionary or list, convert to JSON string
        if isinstance(dialogue, (dict, list)):
            dialogue = json.dumps(dialogue, ensure_ascii=False, indent=4)
            
        # Ensure dialogue is a string
        if not isinstance(dialogue, str):
            dialogue = str(dialogue)
            
        # If dialogue is already formatted JSON string, write directly
        # Otherwise try to parse and reformat
        if not dialogue.startswith('{') and not dialogue.startswith('['):
            # Already text, write directly
            with open(filename, "w", encoding="utf-8") as f:
                f.write(dialogue)
        else:
            # Try to parse and reformat
            try:
                json_obj = json.loads(dialogue)
                with open(filename, "w", encoding="utf-8") as f:
                    json.dump(json_obj, f, ensure_ascii=False, indent=4)
            except:
                # If parsing fails, write directly
                with open(filename, "w", encoding="utf-8") as f:
                    f.write(dialogue)
                    
        return True
    except Exception as e:
        print(f"Error when saving dialogue: {e}")
        return False

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:
        tuple: (List of character profile strings, List of character file paths)
    """
    profiles = []
    expanded_paths = []
    character_files = []  # Store matched file paths
    
    # First 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)
                character_files.append(file_path)
                print(f"Loaded character: {file_path}")
            except Exception as e:
                print(f"Failed to load character file {file_path}: {e}")
    
    return profiles, character_files

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"]

def get_latest_file(pattern: str, default: str = None):
    """Get the most recent file matching a pattern.

    Args:
        pattern (str): Glob pattern to match files
        default (str, optional): Default value if no file is found. Defaults to None.

    Returns:
        str: Path to the most recent file or default value
    """
    files = glob.glob(pattern)
    if not files:
        return default
    return sorted(files)[-1]

def load_selected_characters():
    """Load selected character file paths

    Returns:
        List[str]: List of character file paths, if no selection result is found, return an empty list
    """
    try:
        with open('selected_characters.json', 'r', encoding='utf-8') as f:
            return json.load(f)['selected_files']
    except (FileNotFoundError, json.JSONDecodeError, KeyError):
        return []

if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Generate a dialogue with specified characters')
    parser.add_argument('--scene', type=str, help='Scene file path')
    parser.add_argument('--worldview', type=str, 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')
    parser.add_argument('--output', type=str, help='Dialogue output file path, default is dialogue_{timestamp}.json')
    parser.add_argument('--max-retries', type=int, default=3, help='Maximum number of retries when generating fails')
    parser.add_argument('--use-selector', action='store_true', help='Use character selector tool')
    
    args = parser.parse_args()
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    try:
        # Set default values
        if not args.scene:
            args.scene = get_latest_file("scene/scene_*.yaml", "scene/scene_20250503_213029.yaml")
            print(f"No scene file specified, using latest scene: {args.scene}")
            
        if not args.worldview:
            args.worldview = get_latest_file("worldview/*.yaml", "worldview/worldview_20250503_185119.yaml")
            print(f"No worldview file specified, using latest worldview: {args.worldview}")
        
        # Check if character selector is enabled
        if args.use_selector:
            try:
                print("Starting character selector...")
                from character_selector import launch_selector
                launch_selector()
                
                # Load selected characters
                selected_files = load_selected_characters()
                if selected_files:
                    args.characters = selected_files
                    print(f"Loaded {len(selected_files)} characters from selector")
                else:
                    print("No characters obtained from selector, using default characters")
            except ImportError:
                print("Character selector module not found, please ensure character_selector.py file exists")
            except Exception as e:
                print(f"Failed to start character selector: {e}")
        
        # If no character file is specified and no characters are loaded from selector, use default path
        if not args.characters:
            # Try to load from selector result
            selected_files = load_selected_characters()
            if selected_files:
                args.characters = selected_files
                print(f"Loaded {len(selected_files)} characters from selector")
            else:
                args.characters = get_default_character_paths()
                print(f"No character file specified, using default path: {args.characters}")
            
        if not args.output:
            os.makedirs("dialogue", exist_ok=True)
            args.output = f"dialogue/dialogue_{timestamp}.json"
        
        # Create output directory
        os.makedirs(os.path.dirname(args.output), exist_ok=True)
            
        # Read file
        try:
            scene_content = read_yaml_to_string(args.scene)
            print(f"Successfully read scene file: {args.scene}")
        except Exception as e:
            print(f"Failed to read scene file: {e}")
            print("Using empty scene instead")
            scene_content = ""
            
        try:
            worldview_content = read_yaml_to_string(args.worldview)
            print(f"Successfully read worldview file: {args.worldview}")
        except Exception as e:
            print(f"Failed to read worldview file: {e}")
            print("Using empty worldview instead")
            worldview_content = ""
        
        # Read specified character files
        character_profiles, character_files = load_character_profiles(args.characters)
        
        if not character_profiles:
            print("No valid character files found, cannot generate dialogue")
            exit(1)
            
        print(f"Loaded {len(character_profiles)} characters")
        
        # Extract character names from file names
        character_names = []
        for file_path in character_files:
            # Get file name (without extension)
            base_name = os.path.basename(file_path)
            name_without_ext = os.path.splitext(base_name)[0]
            # Replace underscores with spaces
            character_name = name_without_ext.replace('_', ' ')
            character_names.append(character_name)
        
        print(f"Extracted character names from file names: {', '.join(character_names)}")
        
        # Generate dialogue
        for retry in range(args.max_retries):
            try:
                print(f"Attempting to generate dialogue {retry+1}/{args.max_retries}")
                dialogue = generate_dialogue(
                    scene_content, 
                    worldview_content, 
                    character_profiles, 
                    args.model, 
                    temperature=args.temperature,
                    character_names=character_names
                )
                
                # Check if dialogue contains error
                if isinstance(dialogue, str) and '"error":' in dialogue:
                    print(f"Error when generating dialogue: {dialogue}")
                    if retry < args.max_retries - 1:
                        print("Retrying...")
                        continue
                
                print("\nGenerated dialogue:")
                if isinstance(dialogue, str) and len(dialogue) > 500:
                    print(dialogue[:500] + "...")
                else:
                    print(dialogue)

                # Save dialogue
                if save_dialogue(dialogue, args.output):
                    print(f"Dialogue saved to {args.output}")
                    
                    # Record used files
                    record_file = os.path.splitext(args.output)[0] + "_files.txt"
                    with open(record_file, "w", encoding="utf-8") as f:
                        f.write(f"Scene: {args.scene}\n")
                        f.write(f"Worldview: {args.worldview}\n")
                        f.write("Characters:\n")
                        for path in args.characters:
                            f.write(f"  - {path}\n")
                    print(f"Recorded used files to {record_file}")
                    
                    # Generation successful, exit retry loop
                    break
            except Exception as e:
                print(f"Failed to generate or save dialogue: {e}")
                if retry < args.max_retries - 1:
                    print("Retrying...")
                else:
                    print("Maximum retries reached, giving up")
                    raise
            
    except Exception as e:
        import traceback
        print(f"Error when generating dialogue: {e}")
        traceback.print_exc()