from enum import IntEnum, StrEnum, auto
from typing import Any, Literal, Union

from pydantic import BaseModel

from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from core.tools.entities.tool_entities import ToolSelector
from graphon.entities.base_node_data import BaseNodeData
from graphon.enums import BuiltinNodeTypes, NodeType


class AgentNodeData(BaseNodeData):
    type: NodeType = BuiltinNodeTypes.AGENT
    agent_strategy_provider_name: str
    agent_strategy_name: str
    agent_strategy_label: str
    memory: MemoryConfig | None = None
    # The version of the tool parameter.
    # If this value is None, it indicates this is a previous version
    # and requires using the legacy parameter parsing rules.
    tool_node_version: str | None = None

    class AgentInput(BaseModel):
        value: Union[list[str], list[ToolSelector], Any]
        type: Literal["mixed", "variable", "constant"]

    agent_parameters: dict[str, AgentInput]


class ParamsAutoGenerated(IntEnum):
    CLOSE = 0
    OPEN = 1


class AgentOldVersionModelFeatures(StrEnum):
    """
    Enum class for old SDK version llm feature.
    """

    TOOL_CALL = "tool-call"
    MULTI_TOOL_CALL = "multi-tool-call"
    AGENT_THOUGHT = "agent-thought"
    VISION = auto()
    STREAM_TOOL_CALL = "stream-tool-call"
    DOCUMENT = auto()
    VIDEO = auto()
    AUDIO = auto()