#3 GENERATE CHARACTER PROFILE
import os
import json
from utils import get_llm_response, extract_yaml_from_string, read_yaml_to_string
from prompt import CHARACTER_PROFILE_TEMPLATE
from datetime import datetime
from worldview_generation import generate_worldview_from_scratch

def update_tracking(worldview_path, character_profile_path, character_count, separated_path):
    """
    Update tracking.json, add the relationship between worldview and character profile file
    Args:
        worldview_path: Worldview file path
        character_profile_path: Character profile file path
        character_count: Character count
        separated_path: Separated character file save path
    Returns:
        None
    """
    tracking_path = "./tracking.json"
    if os.path.exists(tracking_path):
        with open(tracking_path, "r", encoding="utf-8") as f:
            tracking_data = json.load(f)
    else:
        tracking_data = []
    
    # Add new record
    new_record = {
        "from_worldview": worldview_path,
        "to_character_profile": character_profile_path,
        "character_count": character_count,
        "separated_path": separated_path
    }
    tracking_data.append(new_record)
    
    # Save updated tracking.json
    with open(tracking_path, "w", encoding="utf-8") as f:
        json.dump(tracking_data, f, indent=4, ensure_ascii=False)

def generate_character_profile(worldview: str, model_name: str, number_of_characters: int = 3, temperature: float = 1.0) -> str:
    """
    Generate character
    Args:
        worldview: Worldview
        number_of_characters: Character count
        model_name: Model name
        temperature: Temperature coefficient
    Returns:
        Character string
    """
    prompt = f"""
    Please create {number_of_characters} expressive characters based on the following worldview, the characters can have close relationships, or can be significantly different, not necessarily humans, even if they cannot speak, do not be limited by the framework of the real world, unleash your imagination, and mark the Optional field can be left blank:
    {worldview}
    Please strictly follow the following yaml format:
    ```yaml
    Character1:
        {CHARACTER_PROFILE_TEMPLATE}
    Character2:
        ...
    Character3:
        ...
    ```
    """
    response = get_llm_response(prompt, model_name, temperature=temperature)
    return extract_yaml_from_string(response)

def save_character_profile(character_profile: str, filename: str) -> None:
    """
    Save character string to file
    Args:
        character_profile: Character string
        filename: File name
    Returns:
        None
    """
    with open(filename, "w", encoding="utf-8") as f:
        f.write(character_profile)

def save_separated_character_profile(character_profile: str, dirname: str) -> None:
    """
    Save the character string of Character1, Character2, Character3 to separate files, using the Full_Name of the character as the file name
    Args:
        character_profile: Character string
        dirname: Save directory name
    Returns:
        None
    """
    # Ensure the directory exists
    os.makedirs(dirname, exist_ok=True)
    
    # Split character information
    characters = {}
    current_character = None
    current_content = []
    
    # Process string by line
    lines = character_profile.strip().split('\n')
    for line in lines:
        # Check if it is a character definition line (e.g. Character1:)
        if line.strip() and line.strip().startswith('Character') and line.strip().endswith(':'):
            # If there is character information, save the previous character
            if current_character:
                characters[current_character] = '\n'.join(current_content)
                current_content = []
            
            # Set the current character name
            current_character = line.strip()[:-1]
        
        # If there is already the current character, add content
        if current_character:
            # Indent processing: remove the first level of indentation (because it is not nested after separation)
            if line.startswith('  '):
                current_content.append(line[2:])
            else:
                current_content.append(line)
    
    # Process the last character
    if current_character and current_content:
        characters[current_character] = '\n'.join(current_content)
    
    # Save each character to a separate file, using the Full_Name attribute of the character as the file name
    for character, content in characters.items():
        # Try to extract the Full_Name attribute from the content
        full_name = None
        name = None  # Backup option
        
        for line in content.split('\n'):
            if line.strip().startswith('Full_Name:'):
                # Extract the Full_Name value, remove quotes
                name_value = line.strip()[10:].strip()
                if name_value.startswith('"') and name_value.endswith('"'):
                    full_name = name_value[1:-1]
                else:
                    full_name = name_value
                break
            elif line.strip().startswith('Name:') and not name:
                # Extract the Name value as a backup, remove quotes
                name_value = line.strip()[5:].strip()
                if name_value.startswith('"') and name_value.endswith('"'):
                    name = name_value[1:-1]
                else:
                    name = name_value
        
        # Select the file name: prioritize Full_Name, then Name, then CharacterX
        file_name = full_name or name or character
        
        # Replace illegal characters in the file name
        safe_name = ''.join(c for c in file_name if c.isalnum() or c in ' .-_').strip()
        safe_name = safe_name.replace(' ', '_')
        
        # If the processed name is empty, use the original character identifier
        if not safe_name:
            safe_name = character
            
        file_path = os.path.join(dirname, f"{safe_name}.yaml")
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content)
        print(f"Saved {character} (Full_Name: {file_name}) to {file_path}")

def generate_character_profile_from_scratch(model_name: str, temperature: float, number_of_characters: int = 3) -> str:
    """
    Generate character from scratch
    Args:
        model_name: Model name
        temperature: Temperature
    Returns:
        character_profile: Character string
    """
    # Initialize timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    print("generating character profile from scratch...")
    print("step 1: selecting best topic and generating worldview...")
    worldview_path = generate_worldview_from_scratch(model_name, temperature)
    
    # Read the generated worldview
    print("step 2: generating character profile from new worldview...")
    worldview = read_yaml_to_string(worldview_path)
    character_profile = generate_character_profile(worldview, model_name, temperature=temperature, number_of_characters=number_of_characters)
    print("generation done")
    print("saving...")
    print(character_profile)
    
    # Create save directory
    os.makedirs("character_profile", exist_ok=True)
    os.makedirs(f"character_pool/separated_{timestamp}", exist_ok=True)
    
    # Save the complete character information to a file
    character_profile_path = f"character_profile/character_profile_{timestamp}.yaml"
    separated_path = f"character_pool/separated_{timestamp}"
    
    save_character_profile(character_profile, character_profile_path)
    
    # Save each character to a separate file
    save_separated_character_profile(character_profile, separated_path)
    
    # Update tracking.json
    update_tracking(worldview_path, character_profile_path, number_of_characters, separated_path)
    
    print(f"saved")
    return character_profile

def generate_character_profile_from_saved_worldview(worldview_path: str, model_name: str, temperature: float, number_of_characters: int = 3) -> str:
    """
    Generate character from saved worldview
    Args:
        worldview_path: Saved worldview file path
        model_name: Model name
        temperature: Temperature
        number_of_characters: Character count
    Returns:
        character_profile: Character string
    """
    # Initialize timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    saved_worldview = read_yaml_to_string(worldview_path)
    # print(saved_worldview)
    print("generating character profile from saved worldview...")
    character_profile = generate_character_profile(saved_worldview, model_name, number_of_characters=number_of_characters, temperature=temperature)
    print("generation done")
    print("saving...")

    os.makedirs("character_profile", exist_ok=True)
    os.makedirs(f"character_pool/separated_{timestamp}", exist_ok=True)
    
    character_profile_path = f"character_profile/character_profile_{timestamp}.yaml"
    separated_path = f"character_pool/separated_{timestamp}"
    
    save_character_profile(character_profile, character_profile_path)
    save_separated_character_profile(character_profile, separated_path)
    
    # Update tracking.json
    update_tracking(worldview_path, character_profile_path, number_of_characters, separated_path)
    
    print(f"saved") 
    return character_profile

if __name__ == "__main__":
    # character_profile = generate_character_profile_from_scratch(model_name="claude-3-7-sonnet-20250219", temperature=1.0)
    # character_profile = generate_character_profile_from_scratch(model_name="gemini-2.5-pro-preview-06-05", temperature=1.0)
    character_profile = generate_character_profile_from_saved_worldview(worldview_path="worldview/worldview_20250623_165609.yaml", model_name="claude-3-7-sonnet-20250219", temperature=1.0)