from __future__ import annotations

from enum import StrEnum
from typing import Any


class MultiModalDecision(StrEnum):
    """基于消息内容分析的选路决策。

    每个值代表所选模型必须满足的一项能力需求。
    PLAIN_TEXT 表示不需要特殊能力。
    """

    PLAIN_TEXT = "plain_text"
    VISION = "vision"
    THINKING = "thinking"
    TOOL_USE = "tool_use"
    STRUCTURED_OUTPUT = "structured_output"


class MultiModalRouter:
    """分析消息列表以确定所需的模型能力。

    用法::

        messages = [{"role": "user", "content": "Hello"}]
        decisions = MultiModalRouter.decide_from_messages(messages)
        # -> {MultiModalDecision.PLAIN_TEXT}

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "What's in this image?"},
                    {"type": "image_url", "image_url": {"url": "..."}},
                ],
            },
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        # -> {MultiModalDecision.VISION}
    """

    @staticmethod
    def decide_from_messages(messages: list[dict[str, Any]]) -> set[MultiModalDecision]:
        """分析消息并返回所需能力的集合。

        参数:
            messages: 消息字典列表，每个字典至少包含 ``role`` 和 ``content`` 键（兼容 OpenAI 格式）。

        返回:
            :class:`MultiModalDecision` 值的集合。如果未检测到特殊能力，
            返回 ``{PLAIN_TEXT}``。
        """
        if not messages:
            return {MultiModalDecision.PLAIN_TEXT}

        decisions: set[MultiModalDecision] = set()

        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")

            # --- VISION: content is a list of blocks with image type ---
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict):
                        block_type = block.get("type", "")
                        if block_type in ("image_url", "image"):
                            decisions.add(MultiModalDecision.VISION)

            # --- TOOL_USE: assistant sent tool_calls, or message is tool result ---
            if msg.get("tool_calls"):
                decisions.add(MultiModalDecision.TOOL_USE)
            if role == "tool":
                decisions.add(MultiModalDecision.TOOL_USE)

            # --- THINKING: assistant included reasoning ---
            if msg.get("reasoning"):
                decisions.add(MultiModalDecision.THINKING)

        return decisions if decisions else {MultiModalDecision.PLAIN_TEXT}