from __future__ import annotations

from enum import Enum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, field_validator

from schemas.agent.ui_hints import UiHintsPayload


class TaskType(str, Enum):
    KNOWLEDGE = "knowledge"
    RECOMMENDATION = "recommendation"
    PLANNING = "planning"
    SCHEDULING = "scheduling"
    REMINDER_MANAGEMENT = "reminder_management"
    TODO_MANAGEMENT = "todo_management"
    COMMUNICATION_DRAFTING = "communication_drafting"
    INFORMATION_ORGANIZATION = "information_organization"
    STATUS_TRACKING = "status_tracking"
    TRANSACTION_ASSIST = "transaction_assist"
    ACTION_EXECUTION = "action_execution"
    TROUBLESHOOTING = "troubleshooting"
    UNKNOWN = "unknown"


class ResultType(str, Enum):
    DIRECT_ANSWER = "direct_answer"
    OPTIONS_WITH_RECOMMENDATION = "options_with_recommendation"
    ACTION_PLAN = "action_plan"
    SCHEDULE_PROPOSAL = "schedule_proposal"
    TODO_LIST = "todo_list"
    DRAFT_MESSAGE = "draft_message"
    SUMMARY = "summary"
    PROGRESS_SUMMARY = "progress_summary"
    DIAGNOSIS_REPORT = "diagnosis_report"
    STRUCTURED_PAYLOAD = "structured_payload"
    EXECUTION_REPORT = "execution_report"
    CLARIFICATION_REQUEST = "clarification_request"
    SAFETY_BLOCK = "safety_block"
    ERROR_REPORT = "error_report"
    UNKNOWN = "unknown"


class TaskTyping(BaseModel):
    model_config = ConfigDict(extra="forbid")

    primary: TaskType
    secondary: list[TaskType] = Field(default_factory=list, max_length=3)


class ResultTyping(BaseModel):
    model_config = ConfigDict(extra="forbid")

    primary: ResultType
    secondary: list[ResultType] = Field(default_factory=list, max_length=3)


class ExecutionMode(str, Enum):
    ONESTEP = "onestep"
    TOOL_ASSISTED = "tool_assisted"
    MULTISTEP = "multistep"


class RunStatus(str, Enum):
    SUCCESS = "success"
    PARTIAL_SUCCESS = "partial_success"
    FAILED = "failed"


class ToolStatus(str, Enum):
    SUCCESS = "success"
    FAILURE = "failure"
    PARTIAL = "partial"


class KeyEntity(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: str
    type: str
    value: str | None = None

    @field_validator("value", mode="before")
    @classmethod
    def normalize_value(cls, value: object) -> object:
        if value is None:
            return None
        if isinstance(value, str):
            return value
        if isinstance(value, bool | int | float):
            return str(value)
        return value


class ConstraintItem(BaseModel):
    model_config = ConfigDict(extra="forbid")

    key: str
    value: str
    required: bool = True

    @field_validator("value", mode="before")
    @classmethod
    def normalize_value(cls, value: object) -> object:
        if isinstance(value, bool | int | float):
            return str(value)
        return value


class NormalizedTaskInput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    user_text: str
    multimodal_summary: list[str] = Field(default_factory=list)
    context_summary: str = Field(default="", max_length=2000)


class RouterAgentOutput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    normalized_task_input: NormalizedTaskInput
    key_entities: list[KeyEntity] = Field(default_factory=list)
    constraints: list[ConstraintItem] = Field(default_factory=list)
    task_typing: TaskTyping
    execution_mode: ExecutionMode
    result_typing: ResultTyping


class ErrorInfo(BaseModel):
    model_config = ConfigDict(extra="forbid")

    code: str
    message: str
    retryable: bool = False
    details: dict[str, Any] | None = None


class ToolAgentOutput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: str
    tool_call_id: str
    tool_call_args: dict[str, Any] | None = None
    status: ToolStatus
    result: str
    error: ErrorInfo | None = None


class WorkerAgentOutputLite(BaseModel):
    model_config = ConfigDict(extra="forbid")

    status: RunStatus = RunStatus.SUCCESS
    answer: str
    key_points: list[str] = Field(default_factory=list)
    result_type: ResultType = ResultType.UNKNOWN
    suggested_actions: list[str] = Field(default_factory=list)
    error: ErrorInfo | None = None


class WorkerAgentOutputRich(WorkerAgentOutputLite):
    ui_hints: UiHintsPayload | None = None


class AgentOutput(WorkerAgentOutputRich):
    model_config = ConfigDict(extra="forbid")


WorkerAgentOutput = WorkerAgentOutputLite | WorkerAgentOutputRich


def resolve_worker_output_model(
    execution_mode: ExecutionMode,
) -> type[WorkerAgentOutputLite]:
    if execution_mode == ExecutionMode.ONESTEP:
        return WorkerAgentOutputLite
    return WorkerAgentOutputRich