agent.py

"""Agent Models."""

from __future__ import annotations

from enum import Enum
from typing import (
    Annotated,
    Any,
    Dict,
    Iterable,
    List,
    Literal,
    Mapping,
    Optional,
    TypeVar,
    Union,
)

from pydantic import (
    BaseModel,
    BeforeValidator,
    ConfigDict,
    Discriminator,
    Field,
    Tag,
    model_validator,
)

from uipath.agent.models._legacy import normalize_legacy_format
from uipath.core.guardrails import (
    BaseGuardrail,
    FieldReference,
    SpecificFieldsSelector,
    UniversalRule,
)
from uipath.eval.mocks import ExampleCall
from uipath.platform.connections import Connection
from uipath.platform.guardrails import (
    BuiltInValidatorGuardrail,
)

EMPTY_SCHEMA = {"type": "object", "properties": {}}


EnumT = TypeVar("EnumT", bound=Enum)


def _match_enum_case_insensitive(enum: type[EnumT], value: str) -> EnumT | None:
    """Find the corresponding enum value, ignoring case."""
    for enum_value in enum:
        if (
            isinstance(enum_value.value, str)
            and enum_value.value.lower() == value.lower()
        ):
            return enum_value
    return None


def _case_insensitive_enum_validator(
    field_name: str,
    enum_cls: type[Enum],
    alias: str | Iterable[str] = (),
) -> BeforeValidator:
    """Create a BeforeValidator for case-insensitive enum normalization."""

    def normalizer(v: Any) -> Any:
        if not isinstance(v, dict):
            return v

        aliases = [alias] if isinstance(alias, str) else list(alias)
        for key in [field_name, *aliases]:
            value = v.get(key)
            if isinstance(value, str):
                match = _match_enum_case_insensitive(enum_cls, value)
                if match is not None:
                    v[key] = match.value
                    break
        return v

    return BeforeValidator(normalizer)


class CaseInsensitiveEnum(Enum):
    """Base class for case-insensitive enums."""

    @classmethod
    def _missing_(cls, value: Any) -> Enum | None:
        """Called when enum value is not found during lookup."""
        if isinstance(value, str):
            match = _match_enum_case_insensitive(cls, value)
            if match is not None:
                return match
        return None


class AgentResourceType(str, CaseInsensitiveEnum):
    """Agent resource type enumeration."""

    TOOL = "tool"
    CONTEXT = "context"
    ESCALATION = "escalation"
    MCP = "mcp"
    UNKNOWN = "unknown"  # fallback branch discriminator


class AgentToolType(str, CaseInsensitiveEnum):
    """Agent tool type enumeration."""

    AGENT = "Agent"
    PROCESS = "Process"
    API = "Api"
    PROCESS_ORCHESTRATION = "ProcessOrchestration"
    INTEGRATION = "Integration"
    INTERNAL = "Internal"
    IXP = "Ixp"
    UNKNOWN = "Unknown"  # fallback branch discriminator


class AgentInternalToolType(str, CaseInsensitiveEnum):
    """Agent internal tool type enumeration."""

    ANALYZE_FILES = "analyze-attachments"
    DEEP_RAG = "deep-rag"
    BATCH_TRANSFORM = "batch-transform"


class AgentEscalationRecipientType(str, CaseInsensitiveEnum):
    """Agent escalation recipient type enumeration."""

    USER_ID = "UserId"
    GROUP_ID = "GroupId"
    USER_EMAIL = "UserEmail"
    ASSET_USER_EMAIL = "AssetUserEmail"
    GROUP_NAME = "GroupName"
    ASSET_GROUP_NAME = "AssetGroupName"
    ARGUMENT_EMAIL = "ArgumentEmail"
    ARGUMENT_GROUP_NAME = "ArgumentGroupName"


class AgentContextRetrievalMode(str, CaseInsensitiveEnum):
    """Agent context retrieval mode enumeration."""

    SEMANTIC = "Semantic"
    STRUCTURED = "Structured"
    DEEP_RAG = "DeepRAG"
    BATCH_TRANSFORM = "BatchTransform"
    UNKNOWN = "Unknown"  # fallback branch discriminator


class AgentMessageRole(str, CaseInsensitiveEnum):
    """Agent message role enumeration."""

    SYSTEM = "system"
    USER = "user"


class AgentGuardrailActionType(str, CaseInsensitiveEnum):
    """Agent guardrail action type enumeration."""

    BLOCK = "block"
    ESCALATE = "escalate"
    FILTER = "filter"
    LOG = "log"
    UNKNOWN = "unknown"  # fallback branch discriminator


class AgentToolArgumentPropertiesVariant(str, CaseInsensitiveEnum):
    """Agent tool argument properties variant enumeration."""

    DYNAMIC = "dynamic"
    ARGUMENT = "argument"
    STATIC = "static"
    TEXT_BUILDER = "textBuilder"


class TextTokenType(str, CaseInsensitiveEnum):
    """Text token type enumeration."""

    SIMPLE_TEXT = "simpleText"
    VARIABLE = "variable"
    EXPRESSION = "expression"


class CitationMode(str, CaseInsensitiveEnum):
    """Citation mode enumeration."""

    INLINE = "Inline"
    SKIP = "Skip"


class DeepRagFileExtension(str, CaseInsensitiveEnum):
    """File extension enumeration for DeepRAG."""

    PDF = "pdf"
    TXT = "txt"


class BatchTransformFileExtension(str, CaseInsensitiveEnum):
    """File extension enumeration for Batch Transform."""

    CSV = "csv"


class BatchTransformWebSearchGrounding(str, CaseInsensitiveEnum):
    """Batch Transform web search grounding enumeration."""

    ENABLED = "Enabled"
    DISABLED = "Disabled"


class BaseCfg(BaseModel):
    """Base configuration model with common settings."""

    model_config = ConfigDict(
        validate_by_name=True, validate_by_alias=True, extra="allow"
    )


class TextToken(BaseCfg):
    """Text token model."""

    type: TextTokenType
    raw_string: str = Field(alias="rawString")


class BaseAgentToolArgumentProperties(BaseCfg):
    """Base tool argument properties model."""

    variant: AgentToolArgumentPropertiesVariant
    is_sensitive: bool = Field(alias="isSensitive")


class AgentToolStaticArgumentProperties(BaseAgentToolArgumentProperties):
    """Static tool argument properties model."""

    variant: Literal[AgentToolArgumentPropertiesVariant.STATIC] = Field(
        default=AgentToolArgumentPropertiesVariant.STATIC, frozen=True
    )
    value: Optional[Any]


class AgentToolArgumentArgumentProperties(BaseAgentToolArgumentProperties):
    """Agent argument argument properties model."""

    variant: Literal[AgentToolArgumentPropertiesVariant.ARGUMENT] = Field(
        default=AgentToolArgumentPropertiesVariant.ARGUMENT,
        frozen=True,
    )
    argument_path: str = Field(alias="argumentPath")


class AgentToolTextBuilderArgumentProperties(BaseAgentToolArgumentProperties):
    """Agent text builder argument properties model."""

    variant: Literal[AgentToolArgumentPropertiesVariant.TEXT_BUILDER] = Field(
        default=AgentToolArgumentPropertiesVariant.TEXT_BUILDER,
        frozen=True,
    )
    tokens: List[TextToken]


AgentToolArgumentProperties = Annotated[
    Union[
        AgentToolStaticArgumentProperties,
        AgentToolArgumentArgumentProperties,
        AgentToolTextBuilderArgumentProperties,
    ],
    Field(discriminator="variant"),
    _case_insensitive_enum_validator("variant", AgentToolArgumentPropertiesVariant),
]


class BaseResourceProperties(BaseCfg):
    """Base resource properties model."""

    example_calls: Optional[list[ExampleCall]] = Field(None, alias="exampleCalls")


class AgentToolSettings(BaseCfg):
    """Agent tool settings model."""

    max_attempts: Optional[int] = Field(None, alias="maxAttempts")
    retry_delay: Optional[int] = Field(None, alias="retryDelay")
    timeout: Optional[int] = Field(None)


class BaseAgentResourceConfig(BaseCfg):
    """Base agent resource configuration model."""

    name: str
    description: str
    is_enabled: bool = Field(default=True, alias="isEnabled")
    # NOTE: this is the union discriminator; don't attach validators here.
    resource_type: Literal[
        AgentResourceType.TOOL,
        AgentResourceType.CONTEXT,
        AgentResourceType.ESCALATION,
        AgentResourceType.MCP,
        AgentResourceType.UNKNOWN,
    ] = Field(alias="$resourceType")


class AgentUnknownResourceConfig(BaseAgentResourceConfig):
    """Fallback for unknown or future resource types."""

    resource_type: Literal[AgentResourceType.UNKNOWN] = Field(
        alias="$resourceType", default=AgentResourceType.UNKNOWN, frozen=True
    )


class AgentContextQuerySetting(BaseCfg):
    """Agent context query setting model."""

    value: str | None = Field(default=None)
    description: str | None = Field(default=None)
    variant: str | None = Field(default=None)


class AgentContextValueSetting(BaseCfg):
    """Agent context value setting model."""

    value: Any = Field(...)


class DeepRagCitationModeSetting(BaseCfg):
    """DeepRAG citation mode setting model."""

    value: CitationMode = Field(...)


class DeepRagFileExtensionSetting(BaseCfg):
    """DeepRAG file extension setting model."""

    value: DeepRagFileExtension = Field(...)


class BatchTransformFileExtensionSetting(BaseCfg):
    """Batch Transform file extension setting model."""

    value: BatchTransformFileExtension = Field(...)


class BatchTransformWebSearchGroundingSetting(BaseCfg):
    """DeepRAG file extension setting model."""

    value: BatchTransformWebSearchGrounding = Field(...)


class AgentContextOutputColumn(BaseCfg):
    """Agent context output column model."""

    name: str = Field(...)
    description: Optional[str] = Field(None)


class AgentContextSettings(BaseCfg):
    """Agent context settings model."""

    result_count: int = Field(alias="resultCount")
    # Allow Unknown explicitly so we can serialize deterministically
    retrieval_mode: AgentContextRetrievalMode = Field(alias="retrievalMode")
    threshold: float = Field(default=0)
    query: Optional[AgentContextQuerySetting] = Field(None)
    folder_path_prefix: Optional[Union[Dict[str, Any], AgentContextValueSetting]] = (
        Field(None, alias="folderPathPrefix")
    )
    file_extension: Optional[Union[Dict[str, Any], AgentContextValueSetting]] = Field(
        None, alias="fileExtension"
    )
    citation_mode: Optional[AgentContextValueSetting] = Field(
        None, alias="citationMode"
    )
    web_search_grounding: Optional[AgentContextValueSetting] = Field(
        None, alias="webSearchGrounding"
    )
    output_columns: Optional[List[AgentContextOutputColumn]] = Field(
        None, alias="outputColumns"
    )


class AgentContextResourceConfig(BaseAgentResourceConfig):
    """Agent context resource configuration model."""

    resource_type: Literal[AgentResourceType.CONTEXT] = Field(
        alias="$resourceType", default=AgentResourceType.CONTEXT, frozen=True
    )
    folder_path: str = Field(alias="folderPath")
    index_name: str = Field(alias="indexName")
    settings: AgentContextSettings = Field(..., description="Context settings")


class AgentMcpTool(BaseCfg):
    """Agent MCP tool model."""

    name: str = Field(..., alias="name")
    description: str = Field(..., alias="description")
    input_schema: Dict[str, Any] = Field(..., alias="inputSchema")
    output_schema: Optional[Dict[str, Any]] = Field(None, alias="outputSchema")
    argument_properties: Dict[str, AgentToolArgumentProperties] = Field(
        {}, alias="argumentProperties"
    )


class DynamicToolsMode(str, CaseInsensitiveEnum):
    """Dynamic tools mode enumeration."""

    NONE = "none"
    SCHEMA = "schema"
    ALL = "all"


class AgentMcpResourceConfig(BaseAgentResourceConfig):
    """Agent MCP resource configuration model."""

    resource_type: Literal[AgentResourceType.MCP] = Field(
        alias="$resourceType", default=AgentResourceType.MCP, frozen=True
    )
    folder_path: str = Field(alias="folderPath")
    slug: str = Field(..., alias="slug")
    available_tools: List[AgentMcpTool] = Field(..., alias="availableTools")
    dynamic_tools: DynamicToolsMode = Field(
        default=DynamicToolsMode.NONE, alias="dynamicTools"
    )


_RECIPIENT_TYPE_NORMALIZED_MAP: Mapping[int | str, AgentEscalationRecipientType] = {
    1: AgentEscalationRecipientType.USER_ID,
    2: AgentEscalationRecipientType.GROUP_ID,
    3: AgentEscalationRecipientType.USER_EMAIL,
    4: AgentEscalationRecipientType.ASSET_USER_EMAIL,
    5: AgentEscalationRecipientType.GROUP_NAME,
    "staticgroupname": AgentEscalationRecipientType.GROUP_NAME,
    6: AgentEscalationRecipientType.ASSET_GROUP_NAME,
    7: AgentEscalationRecipientType.ARGUMENT_EMAIL,
    8: AgentEscalationRecipientType.ARGUMENT_GROUP_NAME,
}


def _normalize_recipient_type(recipient: Any) -> Any:
    """Normalize recipient type from integer or string to enum before discrimination."""
    if not isinstance(recipient, dict):
        return recipient
    recipient_type = recipient.get("type")

    normalized: AgentEscalationRecipientType | None = None
    if isinstance(recipient_type, int):
        normalized = _RECIPIENT_TYPE_NORMALIZED_MAP.get(recipient_type)
    elif isinstance(recipient_type, str):
        normalized = _RECIPIENT_TYPE_NORMALIZED_MAP.get(recipient_type.lower())
        if normalized is None:
            normalized = _match_enum_case_insensitive(
                AgentEscalationRecipientType, recipient_type
            )

    if normalized is not None:
        recipient["type"] = normalized.value

    return recipient


class BaseEscalationRecipient(BaseCfg):
    """Base class for escalation recipients."""

    type: Union[AgentEscalationRecipientType, str] = Field(..., alias="type")


class StandardRecipient(BaseEscalationRecipient):
    """Standard recipient with value field."""

    type: Literal[
        AgentEscalationRecipientType.USER_ID,
        AgentEscalationRecipientType.GROUP_ID,
        AgentEscalationRecipientType.USER_EMAIL,
        AgentEscalationRecipientType.GROUP_NAME,
    ] = Field(..., alias="type")
    value: str = Field(..., alias="value")
    display_name: Optional[str] = Field(default=None, alias="displayName")


class AssetRecipient(BaseEscalationRecipient):
    """Asset recipient with assetName and folderPath."""

    type: Literal[
        AgentEscalationRecipientType.ASSET_USER_EMAIL,
        AgentEscalationRecipientType.ASSET_GROUP_NAME,
    ] = Field(..., alias="type")
    asset_name: str = Field(..., alias="assetName")
    folder_path: str = Field(..., alias="folderPath")


class ArgumentEmailRecipient(BaseEscalationRecipient):
    """Argument email recipient resolved from a named input argument."""

    type: Literal[
        AgentEscalationRecipientType.ARGUMENT_EMAIL,
    ] = Field(..., alias="type")
    argument_name: str = Field(..., alias="argumentName")


class ArgumentGroupNameRecipient(BaseEscalationRecipient):
    """Argument group name recipient resolved from a named input argument."""

    type: Literal[
        AgentEscalationRecipientType.ARGUMENT_GROUP_NAME,
    ] = Field(..., alias="type")
    argument_name: str = Field(..., alias="argumentName")


AgentEscalationRecipient = Annotated[
    Union[StandardRecipient, AssetRecipient, ArgumentEmailRecipient, ArgumentGroupNameRecipient],
    Field(discriminator="type"),
    BeforeValidator(_normalize_recipient_type),
]


class TaskTitleType(str, CaseInsensitiveEnum):
    """Task title type enumeration."""

    DYNAMIC = "dynamic"
    TEXT_BUILDER = "textBuilder"


class BaseTaskTitle(BaseCfg):
    """Base class for task titles."""

    type: Union[TaskTitleType, str] = Field(..., alias="type")


class DynamicTaskTitle(BaseTaskTitle):
    """Dynamic task title with argument path."""

    type: Literal[TaskTitleType.DYNAMIC] = Field(..., alias="type")
    argument_path: str = Field(..., alias="argumentPath")


class TextBuilderTaskTitle(BaseTaskTitle):
    """Text builder task title with tokens."""

    type: Literal[TaskTitleType.TEXT_BUILDER] = Field(..., alias="type")
    tokens: List[TextToken]


TaskTitle = Annotated[
    Union[DynamicTaskTitle, TextBuilderTaskTitle],
    Field(discriminator="type"),
]


def _resolve_task_title(v: Any) -> Any:
    """Resolve taskTitleV2 and taskTitle into a single task_title field."""
    if not isinstance(v, dict):
        return v

    task_title_v2 = v.get("taskTitleV2")
    task_title = v.get("taskTitle")

    # Priority 1: Use taskTitleV2 if present
    if task_title_v2 is not None:
        task_title_type = task_title_v2.get("type")
        if task_title_type is not None:
            normalized_type = _match_enum_case_insensitive(
                TaskTitleType, task_title_type
            )
            if normalized_type is not None:
                task_title_v2["type"] = normalized_type.value
        v["taskTitle"] = task_title_v2
        v.pop("taskTitleV2")
    # Priority 2: Use taskTitle if present (legacy string support)
    elif task_title is not None:
        pass
    else:
        v["taskTitle"] = "Escalation Task"
    return v


class AgentEscalationChannelProperties(BaseResourceProperties):
    """Agent escalation channel properties model."""

    app_name: str | None = Field(default=None, alias="appName")
    app_version: int = Field(..., alias="appVersion")
    folder_name: Optional[str] = Field(None, alias="folderName")
    resource_key: str | None = Field(default=None, alias="resourceKey")
    is_actionable_message_enabled: Optional[bool] = Field(
        None, alias="isActionableMessageEnabled"
    )
    actionable_message_meta_data: Optional[Any] = Field(
        None, alias="actionableMessageMetaData"
    )


class AgentEscalationChannel(BaseCfg):
    """Agent escalation channel model."""

    id: Optional[str] = Field(None, alias="id")
    name: str = Field(..., alias="name")
    type: str = Field(alias="type")
    description: str = Field(..., alias="description")
    input_schema: Dict[str, Any] = Field(..., alias="inputSchema")
    output_schema: Dict[str, Any] = Field(EMPTY_SCHEMA, alias="outputSchema")
    argument_properties: Dict[str, AgentToolArgumentProperties] = Field(
        {}, alias="argumentProperties"
    )
    outcome_mapping: Optional[Dict[str, str]] = Field(None, alias="outcomeMapping")
    properties: AgentEscalationChannelProperties = Field(..., alias="properties")
    recipients: List[AgentEscalationRecipient] = Field(..., alias="recipients")
    task_title: Optional[Union[str, TaskTitle]] = Field(
        default="Escalation Task", alias="taskTitle"
    )
    priority: Optional[str] = None
    labels: List[str] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def _apply_task_title_resolution(cls, v: Any) -> Any:
        """Apply task title resolution."""
        return _resolve_task_title(v)


class AgentEscalationResourceConfig(BaseAgentResourceConfig):
    """Agent escalation resource configuration model."""

    id: Optional[str] = Field(None, alias="id")
    resource_type: Literal[AgentResourceType.ESCALATION] = Field(
        alias="$resourceType", default=AgentResourceType.ESCALATION, frozen=True
    )
    channels: List[AgentEscalationChannel] = Field(alias="channels")
    is_agent_memory_enabled: bool = Field(default=False, alias="isAgentMemoryEnabled")
    escalation_type: Literal[0] = Field(default=0, alias="escalationType")


class AgentIxpVsEscalationProperties(BaseCfg):
    """VS escalation properties model."""

    ixp_tool_id: str = Field(..., alias="ixpToolId")
    storage_bucket_name: str = Field(..., alias="storageBucketName")
    storage_bucket_folder_path: str = Field(..., alias="storageBucketFolderPath")


class AgentIxpVsEscalationResourceConfig(BaseAgentResourceConfig):
    """VS Agent escalation resource configuration model (escalationType=1)."""

    id: Optional[str] = Field(None, alias="id")
    resource_type: Literal[AgentResourceType.ESCALATION] = Field(
        alias="$resourceType", default=AgentResourceType.ESCALATION, frozen=True
    )
    channels: List[AgentEscalationChannel] = Field(alias="channels")
    is_agent_memory_enabled: bool = Field(default=False, alias="isAgentMemoryEnabled")
    escalation_type: Literal[1] = Field(default=1, alias="escalationType")
    vs_escalation_properties: AgentIxpVsEscalationProperties = Field(
        ..., alias="vsEscalationProperties"
    )


class BaseAgentToolResourceConfig(BaseAgentResourceConfig):
    """Base agent tool resource configuration model."""

    resource_type: Literal[AgentResourceType.TOOL] = Field(
        alias="$resourceType", default=AgentResourceType.TOOL, frozen=True
    )
    input_schema: Dict[str, Any] = Field(..., alias="inputSchema")


class AgentProcessToolProperties(BaseResourceProperties):
    """Agent process tool properties model."""

    folder_path: Optional[str] = Field(None, alias="folderPath")
    process_name: Optional[str] = Field(None, alias="processName")


class AgentProcessToolResourceConfig(BaseAgentToolResourceConfig):
    """Agent process tool resource configuration model."""

    type: Literal[
        AgentToolType.AGENT,
        AgentToolType.PROCESS,
        AgentToolType.API,
        AgentToolType.PROCESS_ORCHESTRATION,
    ]
    output_schema: Dict[str, Any] = Field(EMPTY_SCHEMA, alias="outputSchema")
    properties: AgentProcessToolProperties
    settings: AgentToolSettings = Field(default_factory=AgentToolSettings)
    arguments: Dict[str, Any] = Field(default_factory=dict)
    argument_properties: Dict[str, AgentToolArgumentProperties] = Field(
        {}, alias="argumentProperties"
    )


class AgentIxpExtractionToolProperties(BaseResourceProperties):
    """Agent process tool properties model."""

    project_name: str | None = Field(None, alias="projectName")
    version_tag: str | None = Field(None, alias="versionTag")


class AgentIxpExtractionResourceConfig(BaseAgentToolResourceConfig):
    """Agent ixp extraction tool resource configuration model."""

    type: Literal[AgentToolType.IXP] = AgentToolType.IXP
    output_schema: dict[str, Any] = Field(EMPTY_SCHEMA, alias="outputSchema")
    settings: AgentToolSettings = Field(default_factory=AgentToolSettings)
    properties: AgentIxpExtractionToolProperties


class AgentIntegrationToolParameter(BaseCfg):
    """Agent integration tool parameter model."""

    name: str = Field(..., alias="name")
    type: str = Field(..., alias="type")
    value: Optional[Any] = Field(None, alias="value")
    field_location: str = Field(..., alias="fieldLocation")

    # Optional metadata
    display_name: Optional[str] = Field(None, alias="displayName")
    display_value: Optional[str] = Field(None, alias="displayValue")
    description: Optional[str] = Field(None, alias="description")
    position: Optional[str] = Field(None, alias="position")
    field_variant: Optional[str] = Field(None, alias="fieldVariant")
    dynamic: Optional[bool] = Field(None, alias="dynamic")
    is_cascading: Optional[bool] = Field(None, alias="isCascading")
    sort_order: Optional[int] = Field(None, alias="sortOrder")
    required: Optional[bool] = Field(None, alias="required")


class AgentIntegrationToolProperties(BaseResourceProperties):
    """Agent integration tool properties model."""

    tool_path: str = Field(..., alias="toolPath")
    object_name: str = Field(..., alias="objectName")
    tool_display_name: str = Field(..., alias="toolDisplayName")
    tool_description: str = Field(..., alias="toolDescription")
    method: str = Field(..., alias="method")
    connection: Connection = Field(..., alias="connection")
    body_structure: Optional[dict[str, Any]] = Field(None, alias="bodyStructure")
    parameters: List[AgentIntegrationToolParameter] = Field(
        default_factory=list, alias="parameters"
    )


class AgentInternalAnalyzeFilesToolProperties(BaseResourceProperties):
    """Agent internal analyze files tool properties model."""

    tool_type: Literal[AgentInternalToolType.ANALYZE_FILES] = Field(
        alias="toolType", default=AgentInternalToolType.ANALYZE_FILES, frozen=True
    )


class AgentInternalDeepRagToolProperties(BaseResourceProperties):
    """Agent internal DeepRAG tool properties model."""

    tool_type: Literal[AgentInternalToolType.DEEP_RAG] = Field(
        alias="toolType", default=AgentInternalToolType.DEEP_RAG, frozen=True
    )
    settings: AgentInternalDeepRagSettings = Field(..., alias="settings")


class AgentInternalBatchTransformToolProperties(BaseResourceProperties):
    """Agent internal Batch Tranform tool properties model."""

    tool_type: Literal[AgentInternalToolType.BATCH_TRANSFORM] = Field(
        alias="toolType", default=AgentInternalToolType.BATCH_TRANSFORM, frozen=True
    )
    settings: AgentInternalBatchTransformSettings = Field(..., alias="settings")


AgentInternalToolProperties = Annotated[
    Union[
        AgentInternalAnalyzeFilesToolProperties,
        AgentInternalDeepRagToolProperties,
        AgentInternalBatchTransformToolProperties,
    ],
    Field(discriminator="tool_type"),
    _case_insensitive_enum_validator("tool_type", AgentInternalToolType, "toolType"),
]


class AgentInternalDeepRagSettings(BaseCfg):
    """Agent internal DeepRAG tool settings model."""

    context_type: str = Field(..., alias="contextType")
    query: AgentContextQuerySetting = Field(...)
    folder_path_prefix: AgentContextQuerySetting | None = Field(
        default=None, alias="folderPathPrefix"
    )
    citation_mode: DeepRagCitationModeSetting = Field(..., alias="citationMode")
    file_extension: DeepRagFileExtensionSetting = Field(..., alias="fileExtension")


class AgentInternalBatchTransformSettings(BaseCfg):
    """Agent internal Batch Transform tool settings model."""

    context_type: str = Field(..., alias="contextType")
    query: AgentContextQuerySetting = Field(...)
    folder_path_prefix: AgentContextQuerySetting | None = Field(
        default=None, alias="folderPathPrefix"
    )
    file_extension: BatchTransformFileExtensionSetting = Field(
        ..., alias="fileExtension"
    )
    output_columns: List[AgentContextOutputColumn] = Field(..., alias="outputColumns")
    web_search_grounding: BatchTransformWebSearchGroundingSetting = Field(
        ..., alias="webSearchGrounding"
    )


class AgentIntegrationToolResourceConfig(BaseAgentToolResourceConfig):
    """Agent integration tool resource configuration model."""

    type: Literal[AgentToolType.INTEGRATION] = AgentToolType.INTEGRATION
    properties: AgentIntegrationToolProperties
    settings: Optional[AgentToolSettings] = Field(None)
    arguments: Optional[Dict[str, Any]] = Field(default_factory=dict)
    # is output schemas were only recently added so they will be missing in some resources
    output_schema: Optional[Dict[str, Any]] = Field(None, alias="outputSchema")


class AgentInternalToolResourceConfig(BaseAgentToolResourceConfig):
    """Agent internal tool resource configuration model."""

    type: Literal[AgentToolType.INTERNAL] = AgentToolType.INTERNAL
    properties: AgentInternalToolProperties
    settings: Optional[AgentToolSettings] = Field(None)
    arguments: Optional[Dict[str, Any]] = Field(default_factory=dict)
    output_schema: Dict[str, Any] = Field(EMPTY_SCHEMA, alias="outputSchema")
    argument_properties: Dict[str, AgentToolArgumentProperties] = Field(
        {}, alias="argumentProperties"
    )


class AgentUnknownToolResourceConfig(BaseAgentToolResourceConfig):
    """Fallback for unknown tool types (parent normalizer sets type='Unknown')."""

    type: Literal[AgentToolType.UNKNOWN] = AgentToolType.UNKNOWN
    arguments: Optional[Dict[str, Any]] = Field(default_factory=dict)


ToolResourceConfig = Annotated[
    Union[
        AgentProcessToolResourceConfig,
        AgentIntegrationToolResourceConfig,
        AgentInternalToolResourceConfig,
        AgentIxpExtractionResourceConfig,
        AgentUnknownToolResourceConfig,  # when parent sets type="Unknown"
    ],
    Field(discriminator="type"),
]

EscalationResourceConfig = Annotated[
    Union[
        Annotated[AgentEscalationResourceConfig, Tag(0)],
        Annotated[AgentIxpVsEscalationResourceConfig, Tag(1)],
    ],
    Discriminator(lambda v: v.get("escalation_type") or v.get("escalationType") or 0),
]

AgentResourceConfig = Annotated[
    Union[
        ToolResourceConfig,  # nested discrim on 'type'
        AgentContextResourceConfig,
        EscalationResourceConfig,  # nested discrim on 'escalation_type'
        AgentMcpResourceConfig,
        AgentUnknownResourceConfig,  # when parent sets resource_type="Unknown"
    ],
    Field(discriminator="resource_type"),
]


class AgentGuardrailBlockAction(BaseModel):
    """Agent guardrail block action model."""

    action_type: Literal[AgentGuardrailActionType.BLOCK] = Field(
        alias="$actionType", default=AgentGuardrailActionType.BLOCK, frozen=True
    )
    reason: str
    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentGuardrailFilterAction(BaseModel):
    """Agent guardrail filter action model."""

    action_type: Literal[AgentGuardrailActionType.FILTER] = Field(
        alias="$actionType", default=AgentGuardrailActionType.FILTER, frozen=True
    )
    fields: List[FieldReference]
    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentGuardrailSeverityLevel(str, CaseInsensitiveEnum):
    """Severity level enumeration."""

    ERROR = "Error"
    INFO = "Info"
    WARNING = "Warning"


class AgentGuardrailLogAction(BaseModel):
    """Agent guardrail log action model."""

    action_type: Literal[AgentGuardrailActionType.LOG] = Field(
        alias="$actionType", default=AgentGuardrailActionType.LOG, frozen=True
    )
    message: Optional[str] = Field(None, alias="message")
    severity_level: AgentGuardrailSeverityLevel = Field(alias="severityLevel")
    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentGuardrailEscalateActionApp(BaseModel):
    """Agent guardrail escalate action app model."""

    id: Optional[str] = None
    version: int
    name: str
    folder_id: Optional[str] = Field(None, alias="folderId")
    folder_name: str = Field(alias="folderName")
    app_process_key: Optional[str] = Field(None, alias="appProcessKey")
    runtime: Optional[str] = None
    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentGuardrailEscalateAction(BaseModel):
    """Agent guardrail escalate action model."""

    action_type: Literal[AgentGuardrailActionType.ESCALATE] = Field(
        alias="$actionType", default=AgentGuardrailActionType.ESCALATE, frozen=True
    )
    app: AgentGuardrailEscalateActionApp
    recipient: "AgentEscalationRecipient"  # forward ref ok
    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentGuardrailUnknownAction(BaseModel):
    """Fallback for unknown guardrail actions."""

    action_type: Literal[AgentGuardrailActionType.UNKNOWN] = Field(
        alias="$actionType", default=AgentGuardrailActionType.UNKNOWN, frozen=True
    )
    # Accept arbitrary payload for forward-compat
    details: Optional[Dict[str, Any]] = None
    model_config = ConfigDict(populate_by_name=True, extra="allow")


GuardrailAction = Annotated[
    Union[
        AgentGuardrailBlockAction,
        AgentGuardrailFilterAction,
        AgentGuardrailLogAction,
        AgentGuardrailEscalateAction,
        AgentGuardrailUnknownAction,  # when parent sets $actionType="unknown"
    ],
    Field(discriminator="action_type"),
    _case_insensitive_enum_validator(
        "action_type", AgentGuardrailActionType, "$actionType"
    ),
]


class AgentBuiltInValidatorGuardrail(BuiltInValidatorGuardrail):
    """Agent built-in validator guardrail with action capabilities."""

    action: GuardrailAction = Field(
        ..., description="Action to take when guardrail is triggered"
    )

    model_config = ConfigDict(
        validate_by_name=True, validate_by_alias=True, extra="allow"
    )


class AgentWordOperator(str, CaseInsensitiveEnum):
    """Word operator enumeration."""

    CONTAINS = "contains"
    DOES_NOT_CONTAIN = "doesNotContain"
    DOES_NOT_END_WITH = "doesNotEndWith"
    DOES_NOT_EQUAL = "doesNotEqual"
    DOES_NOT_START_WITH = "doesNotStartWith"
    ENDS_WITH = "endsWith"
    EQUALS = "equals"
    IS_EMPTY = "isEmpty"
    IS_NOT_EMPTY = "isNotEmpty"
    MATCHES_REGEX = "matchesRegex"
    STARTS_WITH = "startsWith"


class AgentWordRule(BaseModel):
    """Word rule model."""

    rule_type: Literal["word"] = Field(alias="$ruleType")
    field_selector: AgentFieldSelector = Field(alias="fieldSelector")
    operator: AgentWordOperator
    value: str | None = None

    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentAllFieldsSelector(BaseModel):
    """All fields selector."""

    selector_type: Literal["all"] = Field(alias="$selectorType")

    model_config = ConfigDict(populate_by_name=True, extra="allow")


AgentFieldSelector = Annotated[
    AgentAllFieldsSelector | SpecificFieldsSelector,
    Field(discriminator="selector_type"),
]


class AgentNumberOperator(str, CaseInsensitiveEnum):
    """Number operator enumeration."""

    DOES_NOT_EQUAL = "doesNotEqual"
    EQUALS = "equals"
    GREATER_THAN = "greaterThan"
    GREATER_THAN_OR_EQUAL = "greaterThanOrEqual"
    LESS_THAN = "lessThan"
    LESS_THAN_OR_EQUAL = "lessThanOrEqual"


class AgentNumberRule(BaseModel):
    """Number rule model."""

    rule_type: Literal["number"] = Field(alias="$ruleType")
    field_selector: AgentFieldSelector = Field(alias="fieldSelector")
    operator: AgentNumberOperator
    value: float

    model_config = ConfigDict(populate_by_name=True, extra="allow")


class AgentBooleanOperator(str, CaseInsensitiveEnum):
    """Boolean operator enumeration."""

    EQUALS = "equals"


class AgentBooleanRule(BaseModel):
    """Boolean rule model."""

    rule_type: Literal["boolean"] = Field(alias="$ruleType")
    field_selector: AgentFieldSelector = Field(alias="fieldSelector")
    operator: AgentBooleanOperator
    value: bool

    model_config = ConfigDict(populate_by_name=True, extra="allow")


AgentRule = Annotated[
    AgentWordRule | AgentNumberRule | AgentBooleanRule | UniversalRule,
    Field(discriminator="rule_type"),
]


class AgentCustomGuardrail(BaseGuardrail):
    """Agent custom guardrail with action capabilities."""

    guardrail_type: Literal["custom"] = Field(alias="$guardrailType")
    rules: list[AgentRule]

    action: GuardrailAction = Field(
        ..., description="Action to take when guardrail is triggered"
    )

    model_config = ConfigDict(
        validate_by_name=True, validate_by_alias=True, extra="allow"
    )


class AgentUnknownGuardrail(BaseCfg):
    """Fallback wrapper for unknown guardrail kinds."""

    guardrail_type: Literal["unknown"] = Field(
        alias="$guardrailType", default="unknown", frozen=True
    )
    # store the original payload under 'raw' for round-trip/debug
    raw: Dict[str, Any]


AgentGuardrail = Annotated[
    Union[
        # known kinds from SDK
        AgentCustomGuardrail,
        AgentBuiltInValidatorGuardrail,
        # unknown kind fallback
        AgentUnknownGuardrail,
    ],
    Field(discriminator="guardrail_type"),
]


class AgentMetadata(BaseCfg):
    """Agent metadata model."""

    is_conversational: bool = Field(alias="isConversational")
    storage_version: str = Field(alias="storageVersion")


class AgentMessage(BaseCfg):
    """Agent message model."""

    role: AgentMessageRole
    content: str
    content_tokens: Optional[List[TextToken]] = Field(None, alias="contentTokens")


class AgentByomProperties(BaseCfg):
    """Agent byom properties model."""

    connection_id: str = Field(alias="connectionId")
    connector_key: str = Field(alias="connectorKey")


class AgentSettings(BaseCfg):
    """Agent settings model."""

    engine: str
    model: str
    max_tokens: int = Field(..., alias="maxTokens")
    temperature: float
    byom_properties: Optional[AgentByomProperties] = Field(None, alias="byomProperties")
    max_iterations: Optional[int] = Field(None, alias="maxIterations")


class AgentDefinition(BaseModel):
    """Unified agent definition with parent-level normalization for guardrails and resources."""

    input_schema: Dict[str, Any] = Field(..., alias="inputSchema")
    output_schema: Dict[str, Any] = Field(..., alias="outputSchema")
    guardrails: Optional[List[AgentGuardrail]] = Field(None)

    id: Optional[str] = None
    name: Optional[str] = None
    metadata: Optional[AgentMetadata] = None
    messages: List[AgentMessage]

    version: str = "1.0.0"
    resources: List[AgentResourceConfig] = Field(default_factory=list)
    features: List[Any] = Field(default_factory=list)
    settings: AgentSettings

    model_config = ConfigDict(
        validate_by_name=True, validate_by_alias=True, extra="allow"
    )

    @property
    def is_conversational(self) -> bool:
        """Checks the settings.engine property to determine if the agent is conversational."""
        if hasattr(self, "metadata") and self.metadata:
            metadata = self.metadata
            if hasattr(metadata, "is_conversational"):
                return metadata.is_conversational
        return False

    @staticmethod
    def _normalize_guardrails(v: Dict[str, Any]) -> None:
        guards = v.get("guardrails")
        if not isinstance(guards, list):
            return

        normalized = []
        for g in guards:
            if not isinstance(g, dict):
                normalized.append(g)
                continue

            gt = g.get("$guardrailType") or g.get("guardrail_type")
            # Normalize to the expected discriminator values
            if isinstance(gt, str):
                gt_lower = gt.lower()
                if gt_lower == "custom":
                    g["$guardrailType"] = "custom"
                elif gt_lower == "builtinvalidator":
                    g["$guardrailType"] = "builtInValidator"
                else:
                    # Unknown guardrail type
                    normalized.append({"guardrail_type": "unknown", "raw": g})
                    continue
            else:
                # Non-string guardrail type
                normalized.append({"guardrail_type": "unknown", "raw": g})
                continue

            # Normalize the action if present
            action = g.get("action")
            if isinstance(action, dict):
                at = action.get("$actionType")
                if isinstance(at, str):
                    at_lower = at.lower()
                    if at_lower in {"block", "filter", "log", "escalate"}:
                        # Valid action type, keep as-is or normalize case if needed
                        g["action"]["$actionType"] = at_lower
                    else:
                        # Unknown action type
                        g["action"] = {"$actionType": "unknown", "details": action}
                else:
                    # Non-string action type
                    g["action"] = {"$actionType": "unknown", "details": action}

            normalized.append(g)

        v["guardrails"] = normalized

    @staticmethod
    def _normalize_resources(v: Dict[str, Any]) -> None:
        KNOWN_RES = {"tool", "context", "escalation", "mcp"}
        TOOL_MAP = {
            "agent": "Agent",
            "process": "Process",
            "api": "Api",
            "processorchestration": "ProcessOrchestration",
            "integration": "Integration",
            "internal": "Internal",
            "ixp": "Ixp",
            "unknown": "Unknown",
        }
        CONTEXT_MODE_MAP = {
            "semantic": "Semantic",
            "structured": "Structured",
            "deeprag": "DeepRAG",
            "batchtransform": "BatchTransform",
            "unknown": "Unknown",
        }

        res_list = v.get("resources")
        if not isinstance(res_list, list):
            return

        out = []
        for res in res_list:
            if not isinstance(res, dict):
                out.append(res)
                continue

            rt = res.get("$resourceType") or res.get("resource_type")
            res["$resourceType"] = (
                rt.lower()
                if isinstance(rt, str) and rt.lower() in KNOWN_RES
                else "unknown"
            )

            if res["$resourceType"] == "tool":
                t = res.get("type")
                res["type"] = (
                    TOOL_MAP.get(t.lower(), "Unknown")
                    if isinstance(t, str)
                    else "Unknown"
                )

            if res["$resourceType"] == "context":
                settings = res.get("settings", {})
                rm = settings.get("retrievalMode") or settings.get("retrieval_mode")
                settings["retrievalMode"] = (
                    CONTEXT_MODE_MAP.get(rm.lower(), "Unknown")
                    if isinstance(rm, str)
                    else "Unknown"
                )
                res["settings"] = settings

            out.append(res)

        v["resources"] = out

    @model_validator(mode="before")
    @classmethod
    def _normalize_all(cls, v: Any) -> Any:
        if not isinstance(v, dict):
            return v
        normalize_legacy_format(v)
        cls._normalize_guardrails(v)
        cls._normalize_resources(v)
        return v


LowCodeAgentDefinition = AgentDefinition