"""WebSocket message models for enterprise AI platform integration.

These models capture the bidirectional WebSocket protocol used by enterprise
AI workflow platforms to execute chains/workflows in real-time.

All enterprise-specific data is parameterized -- no hardcoded endpoints,
tokens, or company names.
"""

from __future__ import annotations

import uuid
from datetime import datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field


# --- Enums ---

class WebSocketDirection(StrEnum):
    """Direction of a WebSocket message."""

    SEND = "send"
    RECEIVE = "receive"


class MessageType(StrEnum):
    """Types of WebSocket messages in the protocol."""

    HEARTBEAT = "heartbeat"
    HEARTBEAT_RESPONSE = "heartbeat_response"
    SEND_MESSAGE = "SendMessageV2"
    ANSWER_CHUNK = "answer_chunk"
    TOOL_USE_REQUEST = "tool_use_request"
    TOOL_USE_RESPONSE = "tool_use_response"
    EXECUTION_STATS = "agent_execution_stats"
    AGENT_REASONING = "agent_reasoning"
    ERROR = "error"
    UNKNOWN = "unknown"


class ToolApprovalStatus(StrEnum):
    """Tool execution approval states."""

    APPROVED = "approved"
    PENDING = "pending"
    REJECTED = "rejected"


class ToolExecutionStatus(StrEnum):
    """Tool execution result states."""

    SUCCESS = "success"
    ERROR = "error"
    TIMEOUT = "timeout"
    CANCELLED = "cancelled"


class ExecutionStopReason(StrEnum):
    """Why the agent stopped executing."""

    END_TURN = "end_turn"
    MAX_CYCLES = "max_cycles"
    ERROR = "error"
    USER_CANCELLED = "user_cancelled"


# --- Client → Server Messages ---

class HeartbeatMessage(BaseModel):
    """Client heartbeat to keep WebSocket alive."""

    action: str = Field(default="heartbeat", description="Always 'heartbeat'")


class ModelParam(BaseModel):
    """Configuration for a single component in the workflow."""

    key: str = Field(description="Component parameter key (e.g., 'azure_devops_token')")
    value: str = Field(description="Parameter value (token, setting, etc.)")


class ModelParams(BaseModel):
    """Per-component parameters sent with workflow execution."""

    component_id: str = Field(description="Component identifier (e.g., 'tool_<ID>')")
    params: dict[str, str] = Field(default_factory=dict, description="Key-value parameter pairs")


class InputVariable(BaseModel):
    """A template input variable for the workflow."""

    key: str = Field(description="Variable key (e.g., 'manual_entry_<TS>_<HASH>')")
    label: str = Field(default="", description="Human-readable label")
    value: str = Field(default="", description="User-provided value")
    type: str = Field(default="text", description="Variable type: text, number, select")


class SendMessageV2(BaseModel):
    """Execute a workflow chain via WebSocket.

    This is the primary message sent by the client to trigger
    workflow execution on the enterprise platform.
    """

    action: str = Field(default="SendMessageV2", description="Message action type")
    workflow_id: str = Field(description="UUID of the workflow chain to execute")
    query: str = Field(description="User prompt / instruction for the agent")
    is_persistence_allowed: bool = Field(default=False, description="Whether to persist conversation history")
    modelparams: dict[str, dict[str, str]] = Field(
        default_factory=dict,
        description="Per-component configuration. Keys are component IDs, values are param dicts."
    )
    input_variables: dict[str, str] = Field(
        default_factory=dict,
        description="Template variable values. Keys are variable IDs, values are user inputs."
    )
    conversation_id: str | None = Field(
        default=None,
        description="Previous conversation ID for continuation; null for new conversation"
    )


# --- Server → Client Messages ---

class HeartbeatResponse(BaseModel):
    """Server response to heartbeat."""

    message: str = Field(default="pong", description="Always 'pong'")


class AnswerChunk(BaseModel):
    """A streaming text chunk from the agent's response."""

    agent_id: str = Field(description="Agent node identifier")
    text: str = Field(description="Text fragment of the response")
    timestamp: datetime = Field(default_factory=datetime.now)


class ToolUseRequest(BaseModel):
    """Agent requests to invoke a tool (usually auto-approved)."""

    event_id: str = Field(description="Unique event identifier (e.g., 'toolu_<ID>')")
    event_type: str = Field(default="tool_use", description="Event type")
    tool_use_id: str = Field(description="Tool use tracking ID")
    request_id: str = Field(description="Session-level request UUID")
    tool_name: str = Field(description="Name of the tool to invoke")
    tool_input: dict[str, Any] = Field(default_factory=dict, description="Tool input arguments")
    modified_tool_input: dict[str, Any] | None = Field(default=None, description="Modified input (if guardrails applied)")
    tool_output: Any = Field(default=None, description="Pre-populated output (usually null)")
    tool_status: str | None = Field(default=None, description="Current status")
    approval_status: ToolApprovalStatus = Field(default=ToolApprovalStatus.APPROVED)
    reason: str = Field(default="", description="Approval/rejection reason")
    node_name: str = Field(default="", description="Agent node name in the workflow")
    created_at: datetime = Field(default_factory=datetime.now)
    updated_at: datetime = Field(default_factory=datetime.now)


class ToolUseResponse(BaseModel):
    """Response after a tool has been executed."""

    event_id: str = Field(description="Matching event_id from the request")
    tool_output: list[dict[str, str]] = Field(
        default_factory=list,
        description="Tool output as list of {text: '...'} entries"
    )
    tool_status: ToolExecutionStatus = Field(description="Execution result status")
    tool_exception: str | None = Field(default=None, description="Error details if failed")
    cancel_message: str | None = Field(default=None, description="Cancellation message if cancelled")
    updated_at: datetime = Field(default_factory=datetime.now)


class ToolExecutionStats(BaseModel):
    """Execution statistics for a single tool type."""

    call_count: int = Field(default=0)
    success_count: int = Field(default=0)
    error_count: int = Field(default=0)
    total_time: float = Field(default=0.0, description="Total execution time in seconds")
    average_time: float = Field(default=0.0, description="Average execution time in seconds")
    success_rate: float = Field(default=0.0, description="Success rate 0.0-1.0")


class ToolUsageEntry(BaseModel):
    """Usage statistics for a specific tool."""

    tool_name: str = Field(description="Tool name")
    last_input: dict[str, Any] = Field(default_factory=dict, description="Last input params used")
    stats: ToolExecutionStats = Field(default_factory=ToolExecutionStats)


class AgentExecutionStats(BaseModel):
    """Final execution statistics for a completed agent run."""

    agent_id: str = Field(description="Agent node identifier")
    stop_reason: ExecutionStopReason = Field(description="Why the agent stopped")
    total_cycles: int = Field(default=0, description="Total reasoning cycles")
    total_duration: float = Field(default=0.0, description="Total duration in seconds")
    average_cycle_time: float = Field(default=0.0, description="Average time per cycle in seconds")
    tool_usage: dict[str, ToolUsageEntry] = Field(
        default_factory=dict,
        description="Per-tool execution statistics"
    )


class ReasoningTrace(BaseModel):
    """A single trace entry in the agent's reasoning process."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    name: str = Field(description="Trace name (e.g., 'Cycle 1', 'Tool: get_work_item')")
    parent_id: str | None = Field(default=None, description="Parent trace ID for hierarchy")
    start_time: float = Field(description="Unix timestamp start")
    end_time: float | None = Field(default=None, description="Unix timestamp end")
    duration: float | None = Field(default=None, description="Duration in seconds")
    assistant_text: str | None = Field(default=None, description="Agent's reasoning text for this step")
    children: list[ReasoningTrace] = Field(default_factory=list, description="Child trace entries")


class AgentReasoning(BaseModel):
    """Complete reasoning traces for an agent execution."""

    agent_id: str = Field(description="Agent node identifier")
    traces: list[ReasoningTrace] = Field(default_factory=list)


# --- Parsed WebSocket Message Envelope ---

class ParsedMessage(BaseModel):
    """A parsed WebSocket message with type classification."""

    direction: WebSocketDirection = Field(description="send or receive")
    message_type: MessageType = Field(description="Classified message type")
    agent_id: str | None = Field(default=None, description="Agent ID if present")
    timestamp: datetime = Field(default_factory=datetime.now)
    raw_data: str = Field(default="", description="Original JSON string")

    # Typed payloads (only one will be set based on message_type)
    heartbeat: HeartbeatMessage | None = None
    heartbeat_response: HeartbeatResponse | None = None
    send_message: SendMessageV2 | None = None
    answer_chunk: AnswerChunk | None = None
    tool_use_request: ToolUseRequest | None = None
    tool_use_response: ToolUseResponse | None = None
    execution_stats: AgentExecutionStats | None = None
    agent_reasoning: AgentReasoning | None = None
