# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import enum

from pydantic import BaseModel
from typing import Dict, Any, Optional, Union, List
from enum import Enum

from aworld.config import ConfigDict
from aworld.core.memory import MemoryItem

Config = Union[Dict[str, Any], ConfigDict, BaseModel]


class ActionResult(BaseModel):
    """Result of executing an action by use tool."""
    is_done: bool = False
    success: bool = False
    content: Any = None
    error: str = None
    keep: bool = False
    action_name: str = None
    tool_name: str = None
    # llm tool call id
    tool_call_id: str = None
    metadata: Optional[Dict[str, Any]] = {}
    parameter: Optional[Dict[str, Any]] = {}


class Observation(BaseModel):
    """Observation information is obtained from the tools or transformed from the actions made by agents.

    It can be an agent(as a tool) in the swarm or a tool in the virtual environment.
    """
    # default is None, means the main virtual environment or swarm
    container_id: Optional[str] = None
    # Observer who obtains observation, default is None for compatible, means an agent name or a tool name
    observer: Optional[str] = None
    # default is None for compatible, means with its action/ability name of an agent or a tool
    # NOTE: The only ability of an agent as a tool is handoffs
    ability: Optional[str] = None
    # The agent wants the observation to be created, default is None for compatible.
    from_agent_name: Optional[str] = None
    # To which agent should the observation be given, default is None for compatible.
    to_agent_name: Optional[str] = None
    # general info for agent
    content: Optional[Any] = None
    # dom_tree is a str or DomTree object
    dom_tree: Optional[Union[str, Any]] = None
    image: Optional[str] = None  # base64
    action_result: Optional[List[ActionResult]] = []
    # for video or image list
    images: Optional[List[str]] = []
    # extend key value pair. `done` is an internal key
    info: Optional[Dict[str, Any]] = {}

    @property
    def is_tool_result(self) -> bool:
        return self.action_result is not None and len(self.action_result) > 0


class ParamInfo(BaseModel):
    name: str | None = None
    type: str = "str"
    required: bool = False
    desc: str = None
    default_value: Any = None


class ToolActionInfo(BaseModel):
    name: str
    input_params: Dict[str, ParamInfo] = {}
    desc: str = None


class ActionModel(BaseModel):
    tool_name: Optional[str] = None
    tool_call_id: Optional[str] = None
    # agent name
    agent_name: Optional[str] = None
    # action_name is a tool action name by agent policy.
    action_name: Optional[str] = None
    params: Optional[Dict[str, Any]] = {}
    policy_info: Optional[Any] = None


class TaskItem(BaseModel):
    data: Optional[Any]
    msg: Optional[str] = None
    stop: bool = False
    success: bool = False
    action_name: Optional[str] = None
    params: Optional[Dict[str, Any]] = {}
    policy_info: Optional[Any] = None

class CallbackItem(BaseModel):
    data: Any
    node_id: str = None
    actions: List[ActionModel] = []

class CallbackActionType(str, Enum):
    BYPASS = "bypass"
    OVERRIDE = "override"

class CallbackResult(BaseModel):
    success: bool = False
    result_data: Any = None
    callback_action_type: CallbackActionType = None


class StreamingMode(enum.Enum):
    # core: only core message including agent, tool, chunk, task, group
    CORE = 'core'
    # chunk: only chunk message
    CHUNK = 'chunk'
    # output: only output message
    OUTPUT = 'output'
    # chunk_output: chunk and output message
    CHUNK_OUTPUT = 'chunk_output'
    # all: all message
    ALL = 'all'