from typing import TYPE_CHECKING
from distilabel.steps.base import Step, StepInput

if TYPE_CHECKING:
    from distilabel.steps.typing import StepColumns, StepOutput


class FlipMessageRoles(Step):
    """Flip the roles in chat messages between 'user' and 'assistant'.

    `FlipMessageRoles` is a `Step` that takes a list of chat messages and swaps the roles,
    turning 'user' messages into 'assistant' messages and vice versa. System messages remain unchanged.

    Input columns:
        - messages (`List[Dict[str, str]]`): A list of message dictionaries, where each message has
            'role' and 'content' keys. The role can be 'user', 'assistant', or 'system'.

    Output columns:
        - flipped_messages (`List[Dict[str, str]]`): A list of messages with swapped roles between
            'user' and 'assistant'.

    Categories:
        - format
        - chat
        - role

    Examples:
        Flip roles in chat messages:

        ```python
        from distilabel.steps import FlipMessageRoles

        flip_roles = FlipMessageRoles()
        flip_roles.load()

        result = next(
            flip_roles.process(
                [
                    {
                        "messages": [
                            {"role": "user", "content": "What's 2+2?"},
                            {"role": "assistant", "content": "4"}
                        ]
                    }
                ]
            )
        )
        # >>> result
        # [
        #     {
        #         'messages': [
        #             {'role': 'user', 'content': "What's 2+2?"},
        #             {'role': 'assistant', 'content': '4'}
        #         ],
        #         'flipped_messages': [
        #             {'role': 'assistant', 'content': "What's 2+2?"},
        #             {'role': 'user', 'content': '4'}
        #         ]
        #     }
        # ]
        ```
    """

    @property
    def inputs(self) -> "StepColumns":
        """List of inputs required by the `Step`, which in this case is: `messages`."""
        return ["messages"]

    @property
    def outputs(self) -> "StepColumns":
        """List of outputs generated by the `Step`, which is: `flipped_messages`."""
        return ["flipped_messages"]

    def _flip_role(self, role: str) -> str:
        """Helper method to flip roles between 'user' and 'assistant'.

        Args:
            role: The current role to flip.

        Returns:
            The flipped role, or the original role if it's 'system'.
        """
        role_map = {"user": "assistant", "assistant": "user"}
        return role_map.get(role, role)  # Keeps 'system' role unchanged

    def process(self, *inputs: StepInput) -> "StepOutput":
        """The `process` method flips the roles in the chat messages.

        Args:
            *inputs: A list of `StepInput` containing messages to process.

        Yields:
            A `StepOutput` with batches of `StepInput` with added flipped messages.
        """
        for input in inputs:
            for item in input:
                # Create new messages with flipped roles
                item["flipped_messages"] = [
                    {"role": self._flip_role(msg["role"]), "content": msg["content"]}
                    for msg in item["messages"]
                ]

            yield input
