import re

from openmmla_vision.utils.image_utils import convert_image_to_base64


def generate_vlm_prompt_msg(id_position_dict=None, image_input=None):
    """Generate a prompt message for VLM.

    Args:
        id_position_dict: Dictionary containing the position of each person in the image
        image_input: Either a string (file path) or bytes (image data)

    Returns:
        Formatted prompt message
    """
    if not id_position_dict:
        user_text = (
            "Analyze this laboratory image:\n\n"
            "For each person you can identify in the image:\n"
            "1. Assign a unique ID (e.g., Person 1, Person 2, etc.)\n"
            "2. Describe only what you can see:\n"
            "   a) Approximate position in the image\n"
            "   b) Posture & orientation\n"
            "   c) Gaze direction\n"
            "   d) Hands & objects (if visible)\n"
            "   e) Clear interactions\n"
            "\nFormat:\n"
            "Person [X]:\n"
            "- Position: [brief location]\n"
            "- Visible details: [key observations]\n"
            "\nBe specific about what you can see and what you can't. "
            "If hands or any other part is not visible, state this explicitly. "
            "Do not make assumptions about unseen elements. "
            "Analyze each person you can identify in the image, one by one."
        )
    else:
        user_text = (
            f"Analyze this laboratory image with {len(id_position_dict)} individuals:\n"
        )

        for person_id, position in id_position_dict.items():
            user_text += f"- ID {person_id}: pos {position}\n"

        user_text += (
            "\nFor each person:\n"
            "1. Confirm ID based on position.\n"
            "2. Describe only what you can see:\n"
            "   a) Posture & orientation\n"
            "   b) Gaze direction\n"
            "   c) Hands & objects (if visible)\n"
            "   d) Clear interactions\n"
            "\nFormat:\n"
            "ID [X]:\n"
            "- Position: [brief location]\n"
            "- Visible details: [key observations]\n"
            "\nBe specific about what you can see and what you can't. "
            "If hands or any other part is not visible, state this explicitly. "
            "Do not make assumptions about unseen elements."
        )

    user_text += " Take a deep breath and do the analysis step by step."

    system_content = (
        "You're a capable video ethnographer analyzing a lab experiment."
    )

    if image_input:  # OpenAI API format message
        image_b64 = convert_image_to_base64(image_input)
        vlm_messages = [
            {"role": "system", "content": system_content},
            {"role": "user", "content": [{"type": "text", "text": user_text},
                                         {"type": "image_url", "image_url": {"url": image_b64}}]}
        ]
    else:
        vlm_messages = [
            {"role": "system", "content": system_content},
            {"role": "user", "content": f"(<image>./</image>)\n{user_text}"}
        ]

    return vlm_messages


def generate_llm_prompt_msg(image_description, action_definitions):
    """Generate a prompt message for LLM.

    Args:
        image_description: Description of the image generated by VLM
        action_definitions: Definitions of the actions to categorize

    Returns:
        Formatted prompt message
    """
    user_text = (f"You will now assist with analysing an image description text. Your task is to first categorize each"
                 f"individual's action from the given image description into one of the defined action classes from"
                 f"the action definitions. Then, format the categorization results into a target format.\n\n"
                 f"Image Description:\n{image_description}\n\n"
                 f"Action Definitions:\n{action_definitions}\n\n"
                 f"Carefully read the image description and action definitions and proceed with the following instructions:\n"
                 f"1. Categorize their action into one of the defined classes.\n"
                 f"2. Format your categorization results into a dictionary string with <ID> : <Action> key value pairs,"
                 f" for example, {{'1' : 'Action_1', '7' : 'Action_2'}}.\n"
                 f"Take a deep breath and do the analysis step by step.")

    system_content = (
        "You are now a taxonomist and sociologist who is good at categorizing human behaviors based on the given "
        "context of image description and actions definition. ")

    llm_messages = [
        {"role": "system", "content": system_content},
        {"role": "user", "content": user_text}
    ]

    return llm_messages


def parse_text_to_dict(ai_message: str):
    """Parse a string containing key-value pairs and return a dictionary.

    Args:
        ai_message: The input string to parse.
    Returns:
        A dictionary with keys as ids and values as action descriptions.
    """
    pattern = r"[\"'`]?(\d+)[\"'`]?\s*:\s*[\"'`]?([\w\s]+)[\"'`]?"
    matches = re.finditer(pattern, ai_message)
    return {match.group(1).strip(): match.group(2).strip() for match in matches}
