from collections.abc import Mapping, Sequence
from enum import StrEnum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, JsonValue

from core.app.entities.agent_strategy import AgentStrategyInfo
from core.rag.entities import RetrievalSourceMetadata
from graphon.entities import WorkflowStartReason
from graphon.entities.pause_reason import PauseReasonType
from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from graphon.nodes.human_input.entities import FormInputConfig, UserActionConfig


class AnnotationReplyAccount(BaseModel):
    id: str
    name: str


class AnnotationReply(BaseModel):
    id: str
    account: AnnotationReplyAccount


class TaskStateMetadata(BaseModel):
    annotation_reply: AnnotationReply | None = None
    retriever_resources: Sequence[RetrievalSourceMetadata] = Field(default_factory=list)
    usage: LLMUsage | None = None


class TaskState(BaseModel):
    """
    TaskState entity
    """

    metadata: TaskStateMetadata = Field(default_factory=TaskStateMetadata)


class EasyUITaskState(TaskState):
    """
    EasyUITaskState entity
    """

    llm_result: LLMResult


class WorkflowTaskState(TaskState):
    """
    WorkflowTaskState entity
    """

    answer: str = ""
    first_token_time: float | None = None
    last_token_time: float | None = None
    is_streaming_response: bool = False


class StreamEvent(StrEnum):
    """
    Stream event
    """

    PING = "ping"
    ERROR = "error"
    MESSAGE = "message"
    MESSAGE_END = "message_end"
    TTS_MESSAGE = "tts_message"
    TTS_MESSAGE_END = "tts_message_end"
    MESSAGE_FILE = "message_file"
    MESSAGE_REPLACE = "message_replace"
    AGENT_THOUGHT = "agent_thought"
    AGENT_MESSAGE = "agent_message"
    WORKFLOW_STARTED = "workflow_started"
    WORKFLOW_PAUSED = "workflow_paused"
    WORKFLOW_FINISHED = "workflow_finished"
    NODE_STARTED = "node_started"
    NODE_FINISHED = "node_finished"
    NODE_RETRY = "node_retry"
    ITERATION_STARTED = "iteration_started"
    ITERATION_NEXT = "iteration_next"
    ITERATION_COMPLETED = "iteration_completed"
    LOOP_STARTED = "loop_started"
    LOOP_NEXT = "loop_next"
    LOOP_COMPLETED = "loop_completed"
    TEXT_CHUNK = "text_chunk"
    TEXT_REPLACE = "text_replace"
    AGENT_LOG = "agent_log"
    HUMAN_INPUT_REQUIRED = "human_input_required"
    HUMAN_INPUT_FORM_FILLED = "human_input_form_filled"
    HUMAN_INPUT_FORM_TIMEOUT = "human_input_form_timeout"


class StreamResponse(BaseModel):
    """
    StreamResponse entity
    """

    event: StreamEvent
    task_id: str


class ErrorStreamResponse(StreamResponse):
    """
    ErrorStreamResponse entity
    """

    event: StreamEvent = StreamEvent.ERROR
    err: Exception
    model_config = ConfigDict(arbitrary_types_allowed=True)


class MessageStreamResponse(StreamResponse):
    """
    MessageStreamResponse entity
    """

    event: StreamEvent = StreamEvent.MESSAGE
    id: str
    answer: str
    from_variable_selector: list[str] | None = None


class MessageAudioStreamResponse(StreamResponse):
    """
    MessageStreamResponse entity
    """

    event: StreamEvent = StreamEvent.TTS_MESSAGE
    audio: str


class MessageAudioEndStreamResponse(StreamResponse):
    """
    MessageStreamResponse entity
    """

    event: StreamEvent = StreamEvent.TTS_MESSAGE_END
    audio: str


class MessageEndStreamResponse(StreamResponse):
    """
    MessageEndStreamResponse entity
    """

    event: StreamEvent = StreamEvent.MESSAGE_END
    id: str
    metadata: Mapping[str, object] = Field(default_factory=dict)
    files: Sequence[Mapping[str, Any]] | None = None


class MessageFileStreamResponse(StreamResponse):
    """
    MessageFileStreamResponse entity
    """

    event: StreamEvent = StreamEvent.MESSAGE_FILE
    id: str
    type: str
    belongs_to: str
    url: str


class MessageReplaceStreamResponse(StreamResponse):
    """
    MessageReplaceStreamResponse entity
    """

    event: StreamEvent = StreamEvent.MESSAGE_REPLACE
    answer: str
    reason: str


class AgentThoughtStreamResponse(StreamResponse):
    """
    AgentThoughtStreamResponse entity
    """

    event: StreamEvent = StreamEvent.AGENT_THOUGHT
    id: str
    position: int
    thought: str | None = None
    observation: str | None = None
    tool: str | None = None
    tool_labels: Mapping[str, object] = Field(default_factory=dict)
    tool_input: str | None = None
    message_files: list[str] | None = None


class AgentMessageStreamResponse(StreamResponse):
    """
    AgentMessageStreamResponse entity
    """

    event: StreamEvent = StreamEvent.AGENT_MESSAGE
    id: str
    answer: str


class WorkflowStartStreamResponse(StreamResponse):
    """
    WorkflowStartStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        workflow_id: str
        inputs: Mapping[str, Any]
        created_at: int
        # Always present; mirrors QueueWorkflowStartedEvent.reason for SSE clients.
        reason: WorkflowStartReason = WorkflowStartReason.INITIAL

    event: StreamEvent = StreamEvent.WORKFLOW_STARTED
    workflow_run_id: str
    data: Data


class WorkflowFinishStreamResponse(StreamResponse):
    """
    WorkflowFinishStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        workflow_id: str
        status: WorkflowExecutionStatus
        outputs: Mapping[str, Any] | None = None
        error: str | None = None
        elapsed_time: float
        total_tokens: int
        total_steps: int
        created_by: Mapping[str, object] = Field(default_factory=dict)
        created_at: int
        finished_at: int | None
        exceptions_count: int | None = 0
        files: Sequence[Mapping[str, Any]] | None = []

    event: StreamEvent = StreamEvent.WORKFLOW_FINISHED
    workflow_run_id: str
    data: Data


class WorkflowPauseStreamResponse(StreamResponse):
    """
    WorkflowPauseStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        workflow_run_id: str
        paused_nodes: Sequence[str] = Field(default_factory=list)
        outputs: Mapping[str, Any] = Field(default_factory=dict)
        reasons: Sequence[Mapping[str, Any]] = Field(default_factory=list)
        status: WorkflowExecutionStatus
        created_at: int
        elapsed_time: float
        total_tokens: int
        total_steps: int

    event: StreamEvent = StreamEvent.WORKFLOW_PAUSED
    workflow_run_id: str
    data: Data


class HumanInputRequiredResponse(StreamResponse):
    class Data(BaseModel):
        """
        Data entity
        """

        form_id: str
        node_id: str
        node_title: str
        form_content: str
        inputs: Sequence[FormInputConfig] = Field(default_factory=list)
        actions: Sequence[UserActionConfig] = Field(default_factory=list)
        display_in_ui: bool = False
        form_token: str | None = None
        resolved_default_values: Mapping[str, Any] = Field(default_factory=dict)
        expiration_time: int = Field(..., description="Unix timestamp in seconds")

    event: StreamEvent = StreamEvent.HUMAN_INPUT_REQUIRED
    workflow_run_id: str
    data: Data


class HumanInputRequiredPauseReasonPayload(BaseModel):
    """
    Public pause-reason payload used by blocking responses when only
    ``human_input_required`` events are available.
    """

    TYPE: Literal[PauseReasonType.HUMAN_INPUT_REQUIRED] = PauseReasonType.HUMAN_INPUT_REQUIRED
    form_id: str
    node_id: str
    node_title: str
    form_content: str
    inputs: Sequence[FormInputConfig] = Field(default_factory=list)
    actions: Sequence[UserActionConfig] = Field(default_factory=list)
    display_in_ui: bool = False
    form_token: str | None = None
    resolved_default_values: Mapping[str, Any] = Field(default_factory=dict)
    expiration_time: int

    @classmethod
    def from_response_data(cls, data: HumanInputRequiredResponse.Data) -> "HumanInputRequiredPauseReasonPayload":
        return cls(
            form_id=data.form_id,
            node_id=data.node_id,
            node_title=data.node_title,
            form_content=data.form_content,
            inputs=data.inputs,
            actions=data.actions,
            display_in_ui=data.display_in_ui,
            form_token=data.form_token,
            resolved_default_values=data.resolved_default_values,
            expiration_time=data.expiration_time,
        )


class HumanInputFormFilledResponse(StreamResponse):
    class Data(BaseModel):
        """
        Data entity
        """

        node_id: str
        node_title: str
        rendered_content: str
        action_id: str
        action_text: str

    event: StreamEvent = StreamEvent.HUMAN_INPUT_FORM_FILLED
    workflow_run_id: str
    data: Data


class HumanInputFormTimeoutResponse(StreamResponse):
    class Data(BaseModel):
        """
        Data entity
        """

        node_id: str
        node_title: str
        expiration_time: int

    event: StreamEvent = StreamEvent.HUMAN_INPUT_FORM_TIMEOUT
    workflow_run_id: str
    data: Data


class NodeStartStreamResponse(StreamResponse):
    """
    NodeStartStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        index: int
        predecessor_node_id: str | None = None
        inputs: Mapping[str, Any] | None = None
        inputs_truncated: bool = False
        created_at: int
        extras: dict[str, object] = Field(default_factory=dict)
        iteration_id: str | None = None
        loop_id: str | None = None
        agent_strategy: AgentStrategyInfo | None = None

    event: StreamEvent = StreamEvent.NODE_STARTED
    workflow_run_id: str
    data: Data

    def to_ignore_detail_dict(self) -> dict[str, JsonValue]:
        return {
            "event": self.event.value,
            "task_id": self.task_id,
            "workflow_run_id": self.workflow_run_id,
            "data": {
                "id": self.data.id,
                "node_id": self.data.node_id,
                "node_type": self.data.node_type,
                "title": self.data.title,
                "index": self.data.index,
                "predecessor_node_id": self.data.predecessor_node_id,
                "inputs": None,
                "created_at": self.data.created_at,
                "extras": {},
                "iteration_id": self.data.iteration_id,
                "loop_id": self.data.loop_id,
            },
        }


class NodeFinishStreamResponse(StreamResponse):
    """
    NodeFinishStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        index: int
        predecessor_node_id: str | None = None
        inputs: Mapping[str, Any] | None = None
        inputs_truncated: bool = False
        process_data: Mapping[str, Any] | None = None
        process_data_truncated: bool = False
        outputs: Mapping[str, Any] | None = None
        outputs_truncated: bool = True
        status: WorkflowNodeExecutionStatus
        error: str | None = None
        elapsed_time: float
        execution_metadata: Mapping[WorkflowNodeExecutionMetadataKey, Any] | None = None
        created_at: int
        finished_at: int
        files: Sequence[Mapping[str, Any]] | None = []
        iteration_id: str | None = None
        loop_id: str | None = None

    event: StreamEvent = StreamEvent.NODE_FINISHED
    workflow_run_id: str
    data: Data

    def to_ignore_detail_dict(self) -> dict[str, JsonValue]:
        return {
            "event": self.event.value,
            "task_id": self.task_id,
            "workflow_run_id": self.workflow_run_id,
            "data": {
                "id": self.data.id,
                "node_id": self.data.node_id,
                "node_type": self.data.node_type,
                "title": self.data.title,
                "index": self.data.index,
                "predecessor_node_id": self.data.predecessor_node_id,
                "inputs": None,
                "process_data": None,
                "outputs": None,
                "status": self.data.status,
                "error": None,
                "elapsed_time": self.data.elapsed_time,
                "execution_metadata": None,
                "created_at": self.data.created_at,
                "finished_at": self.data.finished_at,
                "files": [],
                "iteration_id": self.data.iteration_id,
                "loop_id": self.data.loop_id,
            },
        }


class NodeRetryStreamResponse(StreamResponse):
    """
    NodeFinishStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        index: int
        predecessor_node_id: str | None = None
        inputs: Mapping[str, Any] | None = None
        inputs_truncated: bool = False
        process_data: Mapping[str, Any] | None = None
        process_data_truncated: bool = False
        outputs: Mapping[str, Any] | None = None
        outputs_truncated: bool = False
        status: WorkflowNodeExecutionStatus
        error: str | None = None
        elapsed_time: float
        execution_metadata: Mapping[WorkflowNodeExecutionMetadataKey, Any] | None = None
        created_at: int
        finished_at: int
        files: Sequence[Mapping[str, Any]] | None = []
        iteration_id: str | None = None
        loop_id: str | None = None
        retry_index: int = 0

    event: StreamEvent = StreamEvent.NODE_RETRY
    workflow_run_id: str
    data: Data

    def to_ignore_detail_dict(self):
        return {
            "event": self.event.value,
            "task_id": self.task_id,
            "workflow_run_id": self.workflow_run_id,
            "data": {
                "id": self.data.id,
                "node_id": self.data.node_id,
                "node_type": self.data.node_type,
                "title": self.data.title,
                "index": self.data.index,
                "predecessor_node_id": self.data.predecessor_node_id,
                "inputs": None,
                "process_data": None,
                "outputs": None,
                "status": self.data.status,
                "error": None,
                "elapsed_time": self.data.elapsed_time,
                "execution_metadata": None,
                "created_at": self.data.created_at,
                "finished_at": self.data.finished_at,
                "files": [],
                "iteration_id": self.data.iteration_id,
                "loop_id": self.data.loop_id,
                "retry_index": self.data.retry_index,
            },
        }


class IterationNodeStartStreamResponse(StreamResponse):
    """
    NodeStartStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        created_at: int
        extras: dict[str, Any] = Field(default_factory=dict)
        metadata: Mapping = {}
        inputs: Mapping = {}
        inputs_truncated: bool = False

    event: StreamEvent = StreamEvent.ITERATION_STARTED
    workflow_run_id: str
    data: Data


class IterationNodeNextStreamResponse(StreamResponse):
    """
    NodeStartStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        index: int
        created_at: int
        extras: dict[str, Any] = Field(default_factory=dict)

    event: StreamEvent = StreamEvent.ITERATION_NEXT
    workflow_run_id: str
    data: Data


class IterationNodeCompletedStreamResponse(StreamResponse):
    """
    NodeCompletedStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        outputs: Mapping | None = None
        outputs_truncated: bool = False
        created_at: int
        extras: dict[str, Any] | None = None
        inputs: Mapping | None = None
        inputs_truncated: bool = False
        status: WorkflowNodeExecutionStatus
        error: str | None = None
        elapsed_time: float
        total_tokens: int
        execution_metadata: Mapping[str, object] = Field(default_factory=dict)
        finished_at: int
        steps: int

    event: StreamEvent = StreamEvent.ITERATION_COMPLETED
    workflow_run_id: str
    data: Data


class LoopNodeStartStreamResponse(StreamResponse):
    """
    NodeStartStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        created_at: int
        extras: dict[str, Any] = Field(default_factory=dict)
        metadata: Mapping = {}
        inputs: Mapping = {}
        inputs_truncated: bool = False

    event: StreamEvent = StreamEvent.LOOP_STARTED
    workflow_run_id: str
    data: Data


class LoopNodeNextStreamResponse(StreamResponse):
    """
    NodeStartStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        index: int
        created_at: int
        pre_loop_output: Any = None
        extras: Mapping[str, object] = Field(default_factory=dict)

    event: StreamEvent = StreamEvent.LOOP_NEXT
    workflow_run_id: str
    data: Data


class LoopNodeCompletedStreamResponse(StreamResponse):
    """
    NodeCompletedStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        node_id: str
        node_type: str
        title: str
        outputs: Mapping | None = None
        outputs_truncated: bool = False
        created_at: int
        extras: dict[str, Any] | None = None
        inputs: Mapping | None = None
        inputs_truncated: bool = False
        status: WorkflowNodeExecutionStatus
        error: str | None = None
        elapsed_time: float
        total_tokens: int
        execution_metadata: Mapping[str, object] = Field(default_factory=dict)
        finished_at: int
        steps: int

    event: StreamEvent = StreamEvent.LOOP_COMPLETED
    workflow_run_id: str
    data: Data


class TextChunkStreamResponse(StreamResponse):
    """
    TextChunkStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        text: str
        from_variable_selector: list[str] | None = None

    event: StreamEvent = StreamEvent.TEXT_CHUNK
    data: Data


class TextReplaceStreamResponse(StreamResponse):
    """
    TextReplaceStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        text: str

    event: StreamEvent = StreamEvent.TEXT_REPLACE
    data: Data


class PingStreamResponse(StreamResponse):
    """
    PingStreamResponse entity
    """

    event: StreamEvent = StreamEvent.PING


class AppStreamResponse(BaseModel):
    """
    AppStreamResponse entity
    """

    stream_response: StreamResponse


class ChatbotAppStreamResponse(AppStreamResponse):
    """
    ChatbotAppStreamResponse entity
    """

    conversation_id: str
    message_id: str
    created_at: int


class CompletionAppStreamResponse(AppStreamResponse):
    """
    CompletionAppStreamResponse entity
    """

    message_id: str
    created_at: int


class WorkflowAppStreamResponse(AppStreamResponse):
    """
    WorkflowAppStreamResponse entity
    """

    workflow_run_id: str | None = None


class AppBlockingResponse(BaseModel):
    """
    AppBlockingResponse entity
    """

    task_id: str


class ChatbotAppBlockingResponse(AppBlockingResponse):
    """
    ChatbotAppBlockingResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        mode: str
        conversation_id: str
        message_id: str
        answer: str
        metadata: Mapping[str, object] = Field(default_factory=dict)
        created_at: int

    data: Data


class AdvancedChatPausedBlockingResponse(AppBlockingResponse):
    """
    ChatbotAppPausedBlockingResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        mode: str
        conversation_id: str
        message_id: str
        workflow_run_id: str
        answer: str
        metadata: Mapping[str, object] = Field(default_factory=dict)
        created_at: int
        paused_nodes: Sequence[str] = Field(default_factory=list)
        reasons: Sequence[Mapping[str, Any]] = Field(default_factory=list[Mapping[str, Any]])
        status: WorkflowExecutionStatus
        elapsed_time: float
        total_tokens: int
        total_steps: int

    data: Data


class CompletionAppBlockingResponse(AppBlockingResponse):
    """
    CompletionAppBlockingResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        mode: str
        message_id: str
        answer: str
        metadata: Mapping[str, object] = Field(default_factory=dict)
        created_at: int

    data: Data


class WorkflowAppBlockingResponse(AppBlockingResponse):
    """
    WorkflowAppBlockingResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        workflow_id: str
        status: WorkflowExecutionStatus
        outputs: Mapping[str, Any] | None = None
        error: str | None = None
        elapsed_time: float
        total_tokens: int
        total_steps: int
        created_at: int
        finished_at: int | None

    workflow_run_id: str
    data: Data


class WorkflowAppPausedBlockingResponse(AppBlockingResponse):
    """
    WorkflowAppPausedBlockingResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        id: str
        workflow_id: str
        status: WorkflowExecutionStatus
        outputs: Mapping[str, Any] | None = None
        error: str | None = None
        elapsed_time: float
        total_tokens: int
        total_steps: int
        created_at: int
        finished_at: int | None
        paused_nodes: Sequence[str] = Field(default_factory=list)
        reasons: Sequence[Mapping[str, Any]] = Field(default_factory=list)

    workflow_run_id: str
    data: Data


class AgentLogStreamResponse(StreamResponse):
    """
    AgentLogStreamResponse entity
    """

    class Data(BaseModel):
        """
        Data entity
        """

        node_execution_id: str
        id: str
        label: str
        parent_id: str | None = None
        error: str | None = None
        status: str
        data: Mapping[str, Any]
        metadata: Mapping[str, object] = Field(default_factory=dict)
        node_id: str

    event: StreamEvent = StreamEvent.AGENT_LOG
    data: Data