# Copyright (c) Microsoft. All rights reserved.

from typing import Any

from agent_framework._types import Content, Message
from loguru import logger


def _get_role_value(role: Any) -> str:
    """Get the string value of a role, handling both enum and string."""
    return role.value if hasattr(role, "value") else str(role)


def flip_messages(messages: list[Message]) -> list[Message]:
    """Flip message roles between assistant and user for role-playing scenarios.

    Used in agent simulations where the assistant's messages become user inputs
    and vice versa. Function calls are filtered out when flipping assistant
    messages to user messages (since users typically don't make function calls).
    """

    def filter_out_function_calls(messages: list[Content]) -> list[Content]:
        """Remove function call content from message contents."""
        return [content for content in messages if content.type != "function_call"]

    flipped_messages: list[Message] = []
    for msg in messages:
        role_value = _get_role_value(msg.role)
        if role_value == "assistant":
            # Flip assistant to user
            contents = filter_out_function_calls(msg.contents)
            if contents:
                flipped_msg = Message(
                    role="user",
                    # The function calls will cause 400 when role is user
                    contents=contents,
                    author_name=msg.author_name,
                    message_id=msg.message_id,
                )
                flipped_messages.append(flipped_msg)
        elif role_value == "user":
            # Flip user to assistant
            flipped_msg = Message(
                role="assistant", contents=msg.contents, author_name=msg.author_name, message_id=msg.message_id
            )
            flipped_messages.append(flipped_msg)
        elif role_value == "tool":
            # Skip tool messages
            pass
        else:
            # Keep other roles as-is (system, tool, etc.)
            flipped_messages.append(msg)
    return flipped_messages


def log_messages(messages: list[Message]) -> None:
    """Log messages with colored output based on role and content type.

    Provides visual debugging by color-coding different message roles and
    content types. Escapes HTML-like characters to prevent log formatting issues.
    """
    logger_ = logger.opt(colors=True)
    for msg in messages:
        role_value = _get_role_value(msg.role)
        # Handle different content types
        if hasattr(msg, "contents") and msg.contents:
            for content in msg.contents:
                if hasattr(content, "type"):
                    if content.type == "text":
                        escape_text = content.text.replace("<", r"\<")  # type: ignore[union-attr]
                        if role_value == "system":
                            logger_.info(f"<cyan>[SYSTEM]</cyan> {escape_text}")
                        elif role_value == "user":
                            logger_.info(f"<green>[USER]</green> {escape_text}")
                        elif role_value == "assistant":
                            logger_.info(f"<blue>[ASSISTANT]</blue> {escape_text}")
                        elif role_value == "tool":
                            logger_.info(f"<yellow>[TOOL]</yellow> {escape_text}")
                        else:
                            logger_.info(f"<magenta>[{role_value.upper()}]</magenta> {escape_text}")
                    elif content.type == "function_call":
                        function_call_text = f"{content.name}({content.arguments})"
                        function_call_text = function_call_text.replace("<", r"\<")
                        logger_.info(f"<yellow>[TOOL_CALL]</yellow> 🔧 {function_call_text}")
                    elif content.type == "function_result":
                        function_result_text = f"ID:{content.call_id} -> {content.result}"
                        function_result_text = function_result_text.replace("<", r"\<")
                        logger_.info(f"<yellow>[TOOL_RESULT]</yellow> 🔨 {function_result_text}")
                    else:
                        content_text = str(content).replace("<", r"\<")
                        logger_.info(f"<magenta>[{role_value.upper()}] ({content.type})</magenta> {content_text}")
                else:
                    # Fallback for content without type
                    text_content = str(content).replace("<", r"\<")
                    if role_value == "system":
                        logger_.info(f"<cyan>[SYSTEM]</cyan> {text_content}")
                    elif role_value == "user":
                        logger_.info(f"<green>[USER]</green> {text_content}")
                    elif role_value == "assistant":
                        logger_.info(f"<blue>[ASSISTANT]</blue> {text_content}")
                    elif role_value == "tool":
                        logger_.info(f"<yellow>[TOOL]</yellow> {text_content}")
                    else:
                        logger_.info(f"<magenta>[{role_value.upper()}]</magenta> {text_content}")
        elif hasattr(msg, "text") and msg.text:
            # Handle simple text messages
            text_content = msg.text.replace("<", r"\<")
            if role_value == "system":
                logger_.info(f"<cyan>[SYSTEM]</cyan> {text_content}")
            elif role_value == "user":
                logger_.info(f"<green>[USER]</green> {text_content}")
            elif role_value == "assistant":
                logger_.info(f"<blue>[ASSISTANT]</blue> {text_content}")
            elif role_value == "tool":
                logger_.info(f"<yellow>[TOOL]</yellow> {text_content}")
            else:
                logger_.info(f"<magenta>[{role_value.upper()}]</magenta> {text_content}")
        else:
            # Fallback for other message formats
            text_content = str(msg).replace("<", r"\<")
            logger_.info(f"<magenta>[{role_value.upper()}]</magenta> {text_content}")
