import re
import io
import zipfile
import json

def parse_dialogue_tags(dialogue):
    """
    Parses the dialogues with <look_at:character> tags
    and returns clean dialogue and tags
    """
    parts = re.split(r'(<look_at:[a-zA-Z]+>)', dialogue)
    current_word_idx = 0
    clean_text = ""
    tags = []

    for part in parts:
        if part.startswith("<look_at:"):
            # It is a look_at tag
            character = part[9:-1].lower()
            tags.append({
                "character": character.lower(),
                "word_idx": current_word_idx
            })
        else:
            # It is normal part of sentence.
            words = part.strip().split()
            current_word_idx += len(words)
            clean_text += part + " "
    
    return clean_text.strip(), tags

def create_episode_zip(episode_payload, audio_chunks):
    """
    Args:
        episode_payload: Json containing the dialogues
        audio_chunks: dict containing file names mapped to np.float32 arrays
    """
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file:
        zip_file.writestr("episode.json", json.dumps(episode_payload))
        for filename, audio_array in audio_chunks.items():
            zip_file.writestr(filename, audio_array.tobytes())
    
    return zip_buffer.getvalue()

def postprocess_script(script, character_names: dict[str, str]):
    """
    Replaces all named characters with their roles, so unity can use them.
    Args:
        script: The parsed JSON script
        character_names: A dict like {"therapist": "emily", "male": "mike", "female": "mia"}
        
    """
    # Now it's {"emily": "therapist", "mike": "male", "mia": "female"}
    name_to_role = {name.lower(): role for role, name in character_names.items()}

    for dialogue in script.get("dialogues", []):
        char_name = dialogue.get("character", "").lower()
        if char_name in name_to_role:
            dialogue["character"] = name_to_role[char_name]
        
        for look_target in dialogue.get("lookAt", []):
            target_name = look_target.get("character", "").lower()
            if target_name in name_to_role:
                look_target["character"] = name_to_role[target_name]
        
    return script



expected_schema = {
    "type": "object",
    "properties": {
        "dialogues": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "character": {"type": "string"},
                    "dialogue": {"type": "string"},
                    "expression": {"type": "string"}
                },
                "required": ["character", "dialogue", "expression"]
            }
        }
    },
    "required": ["dialogues"]
}

therapistCostumes = [
    "therapist_beauty",
    "therapist_suit",
    "therapist_futuristic",
    "therapist_vampire",
    "therapist_huggingface",
]

maleCostumes = [
    "male_charisma",
    "male_suit",
    "male_shirt",
    "male_vampire",
    "male_hoodie",
]

femaleCostumes = [
    "female_charisma",
    "female_suit",
    "female_futuristic",
    "female_vampire",
    "female_offgrid",
]