"""
Pendo Agent Event Schema

Formal Pydantic models that define the wire format for events sent from the
Python SDK to Pendo's ingestion pipeline.

Aligned with the three-layer schema defined in the web SDK
(pendo-agent-web-sdk-demo) and Go ingestion backend (tasks/newevent.go),
adapted for server-side use:

  Layer 1 — SDK Event (wire format):  AgentEvent / EventBatch
  Layer 2 — ObservationRow:           What ClickHouse stores per span
  Layer 3 — ConversationRow:          Pre-aggregated per conversation (PM-facing)

This module defines Layer 1. Layers 2 & 3 live in the backend.

Server-side adaptations:
  - No web-specific fields (tabId, frameId, sessionId, recordingId, etc.)
  - Timestamps in epoch milliseconds (not nanoseconds)
  - Plain JSON POST to /data/agentic (no JZB encoding)
  - visitor_id and account_id on every event
  - Props carry span-type-specific attributes (LLM model/tokens, tool I/O)

The exporter converts OTel ReadableSpan objects into these models before
serialization and network export.
"""

from __future__ import annotations

import time
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field, field_validator


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class SpanType(str, Enum):
    """The kind of work a span represents.

    Maps to ``pendo.span.type`` on OTel span attributes and the ``type``
    field in Pendo's ingestion schema. Aligned with the web SDK's type enum.
    """

    AGENT = "AGENT"
    GENERATION = "GENERATION"
    TOOL = "TOOL"
    CHAIN = "CHAIN"
    RETRIEVAL = "RETRIEVAL"
    GUARDRAIL = "GUARDRAIL"
    USER_MESSAGE = "USER_MESSAGE"
    SIGNAL = "SIGNAL"
    REACTION = "REACTION"
    UNKNOWN = "UNKNOWN"


class EventType(str, Enum):
    """High-level event classification sent on the wire.

    Distinguishes the *event kind* from the *span type*. A single span may
    emit multiple events (e.g. span_start + span_end).
    """

    MESSAGE = "message"
    TOOL_CALL = "tool_call"
    SIGNAL = "signal"
    REACTION = "reaction"
    SPAN_START = "span_start"
    SPAN_END = "span_end"
    CONVERSATION_START = "conversation_start"
    CONVERSATION_END = "conversation_end"


class SpanStatusCode(str, Enum):
    """Span completion status, mirroring OTel StatusCode."""

    UNSET = "UNSET"
    OK = "OK"
    ERROR = "ERROR"


class ToolStatus(str, Enum):
    """Outcome of a tool/function call."""

    SUCCESS = "success"
    ERROR = "error"


class ReactionType(str, Enum):
    """User feedback reaction types."""

    THUMBS_UP = "thumbs_up"
    THUMBS_DOWN = "thumbs_down"
    COPY = "copy"
    RETRY = "retry"
    EDIT = "edit"


class ConversationOutcome(str, Enum):
    """How a conversation ended."""

    COMPLETED = "completed"
    ABANDONED = "abandoned"
    ESCALATED = "escalated"
    ERROR = "error"


class ContextType(str, Enum):
    """Type of context/file uploaded with a message."""

    CODE = "code"
    PPT = "ppt"
    XL = "xl"
    SHEET = "sheet"
    PDF = "pdf"
    IMAGE = "image"
    TEXT = "text"
    OTHER = "other"


# ---------------------------------------------------------------------------
# Context item — replaces simple fileUploaded boolean
# ---------------------------------------------------------------------------

class ContextItem(BaseModel):
    """A file or context item attached to a message.

    Replaces the old fileUploaded boolean. If any context items are present,
    fileUploaded is auto-set to True for backward compatibility.
    """

    type: ContextType = Field(
        ...,
        description="Type of context: code, ppt, xl, sheet, pdf, image, text, other.",
    )
    title: Optional[str] = Field(
        None,
        description="File name or description of the context item.",
    )


# ---------------------------------------------------------------------------
# Props — span-type-specific properties
# ---------------------------------------------------------------------------

class EventProps(BaseModel):
    """Union of all span-type-specific properties.

    Only fields relevant to the span type will be populated; the rest remain
    None and are excluded from serialization via ``to_wire()``.
    """

    # --- GENERATION (LLM call) fields ---
    model: Optional[List[str]] = Field(
        None,
        description="Model identifier(s). Supports multiple models per message.",
    )
    prompt_tokens: Optional[int] = Field(
        None,
        description="Number of prompt/input tokens consumed.",
        ge=0,
    )
    completion_tokens: Optional[int] = Field(
        None,
        description="Number of completion/output tokens generated.",
        ge=0,
    )
    total_tokens: Optional[int] = Field(
        None,
        description="Total tokens (prompt + completion). Computed if not set.",
        ge=0,
    )
    cost_usd: Optional[float] = Field(
        None,
        description="Estimated cost in USD for this LLM call.",
        ge=0,
    )

    # --- TOOL fields ---
    tool_name: Optional[str] = Field(
        None,
        description="Name of the tool or function that was invoked.",
    )
    tool_input: Optional[str] = Field(
        None,
        description="Stringified input arguments to the tool.",
    )
    tool_output: Optional[str] = Field(
        None,
        description="Stringified return value from the tool.",
    )
    tool_status: Optional[ToolStatus] = Field(
        None,
        description="Whether the tool call succeeded or errored.",
    )

    # --- MESSAGE fields ---
    role: Optional[str] = Field(
        None,
        description="Message role: 'user', 'assistant', or 'system'.",
    )
    content: Optional[str] = Field(
        None,
        description="Message text content.",
    )

    # --- SIGNAL fields ---
    signal_name: Optional[str] = Field(
        None,
        description="Behavioral signal name, e.g. 'user_frustration', 'agent_loop'.",
    )
    signal_data: Optional[Dict[str, Any]] = Field(
        None,
        description="Arbitrary data associated with the signal.",
    )

    # --- REACTION fields ---
    reaction_type: Optional[ReactionType] = Field(
        None,
        description="Type of user feedback reaction.",
    )
    feedback_comment: Optional[str] = Field(
        None,
        description="Written feedback text accompanying a thumbs up/down reaction.",
    )
    message_span_id: Optional[str] = Field(
        None,
        description="Span ID of the message this reaction targets.",
    )

    # --- CONTEXT fields (replaces fileUploaded) ---
    context_uploaded: Optional[List[ContextItem]] = Field(
        None,
        description="List of files/context items attached to this message. "
        "Replaces fileUploaded boolean. If present, fileUploaded is auto-set to true.",
    )
    file_uploaded: Optional[bool] = Field(
        None,
        description="Backward-compat: true if any context items are attached. "
        "Auto-set when context_uploaded is provided.",
    )
    suggested_prompt: Optional[bool] = Field(
        None,
        description="Whether this prompt was a suggested/pre-built prompt.",
    )

    # --- Generic catch-all ---
    extra: Dict[str, Any] = Field(
        default_factory=dict,
        description="Catch-all for custom span attributes not covered above.",
    )


# ---------------------------------------------------------------------------
# Core event model (Layer 1 — Wire Format)
# ---------------------------------------------------------------------------

class AgentEvent(BaseModel):
    """A single event sent to Pendo's ``/data/agentic`` ingestion endpoint.

    This is the server-side equivalent of the JZB-encoded events the web SDK
    sends. Web-only fields (tabId, frameId, sessionId, recordingId,
    recordingSessionId, parentTabId) are intentionally excluded.

    Matches the Layer 1 schema from the web SDK with field names aligned to
    what the Go ingestion backend (``HandleNewEvent`` / ``ParseEvents``)
    expects.
    """

    # --- Identity (required on every event) ---
    visitor_id: str = Field(
        ...,
        description="Pendo visitor ID. Identifies the end-user.",
    )
    account_id: Optional[str] = Field(
        None,
        description="Pendo account ID. Identifies the customer/tenant.",
    )

    # --- Trace context ---
    trace_id: str = Field(
        ...,
        description="32-char lowercase hex trace ID linking related spans.",
    )
    span_id: str = Field(
        ...,
        description="16-char lowercase hex span ID for this event.",
    )
    parent_span_id: Optional[str] = Field(
        None,
        description="16-char lowercase hex span ID of the parent span.",
    )
    depth: int = Field(
        0,
        description="Nesting depth in the span tree. Root = 0.",
        ge=0,
    )

    # --- Conversation context ---
    conversation_id: Optional[str] = Field(
        None,
        description="Groups related spans into a single conversation/session.",
    )

    # --- Event classification ---
    event_type: Optional[EventType] = Field(
        None,
        description="High-level event classification (message, tool_call, signal, etc.).",
    )
    type: SpanType = Field(
        ...,
        description="The kind of work this span represents.",
    )
    name: str = Field(
        ...,
        description="Human-readable span name, e.g. 'classify_intent' or 'search_db'.",
    )
    status: SpanStatusCode = Field(
        SpanStatusCode.UNSET,
        description="Span completion status.",
    )
    error_message: Optional[str] = Field(
        None,
        description="Error description when status is ERROR.",
    )

    # --- Timestamps (epoch milliseconds) ---
    browser_time: int = Field(
        default_factory=lambda: int(time.time() * 1000),
        description="Event timestamp in epoch milliseconds. Named browser_time "
        "for compatibility with Pendo's ingestion pipeline.",
    )
    start_time: int = Field(
        ...,
        description="Span start time in epoch milliseconds.",
    )
    end_time: Optional[int] = Field(
        None,
        description="Span end time in epoch milliseconds.",
    )
    duration_ms: Optional[float] = Field(
        None,
        description="Span duration in milliseconds.",
        ge=0,
    )

    # --- Properties ---
    props: EventProps = Field(
        default_factory=EventProps,
        description="Span-type-specific properties (LLM tokens, tool I/O, etc.).",
    )

    # --- Custom tags ---
    tags: Optional[Dict[str, str]] = Field(
        None,
        description="Custom key-value tags for slicing analytics. "
        "Primary use case: experiment identifiers, e.g. "
        '{"experiment": "new_prompt_v2", "variant": "B"}.',
    )

    # --- SDK metadata ---
    sdk_name: str = Field(
        "pendo_sdk",
        description="Name of the SDK that produced this event.",
    )
    sdk_version: str = Field(
        "0.1.0",
        description="Version of the SDK that produced this event.",
    )

    # --- Validators ---

    @field_validator("trace_id")
    @classmethod
    def _validate_trace_id(cls, v: str) -> str:
        if len(v) != 32 or not all(c in "0123456789abcdef" for c in v):
            raise ValueError("trace_id must be a 32-char lowercase hex string")
        return v

    @field_validator("span_id", "parent_span_id")
    @classmethod
    def _validate_span_id(cls, v: Optional[str]) -> Optional[str]:
        if v is not None:
            if len(v) != 16 or not all(c in "0123456789abcdef" for c in v):
                raise ValueError("span_id must be a 16-char lowercase hex string")
        return v

    def to_wire(self) -> Dict[str, Any]:
        """Serialize to the dict format expected by the ingestion endpoint.

        Excludes None values to keep payloads compact. Flattens ``extra``
        into ``props`` so custom attributes appear as top-level prop keys.
        Handles backward compatibility for modelUsed and fileUploaded.
        """
        data = self.model_dump(exclude_none=True)
        props = data.get("props", {})

        # Flatten extra into props
        extra = props.pop("extra", {})
        props.update(extra)

        # model: wire format uses "modelUsed" — send as list for multi-model support
        # but also keep backward compat single string as "modelUsed" for existing pipeline
        if "model" in props:
            models = props.pop("model")
            if isinstance(models, list) and len(models) == 1:
                props["modelUsed"] = models[0]
                props["modelsUsed"] = models
            elif isinstance(models, list):
                props["modelUsed"] = models[0] if models else ""
                props["modelsUsed"] = models
            else:
                props["modelUsed"] = models
                props["modelsUsed"] = [models]

        # context_uploaded → auto-set fileUploaded for backward compat
        if "context_uploaded" in props:
            props["fileUploaded"] = True
            props["contextUploaded"] = props.pop("context_uploaded")
        elif "file_uploaded" in props:
            props["fileUploaded"] = props.pop("file_uploaded")

        # feedback_comment → feedbackComment wire name
        if "feedback_comment" in props:
            props["feedbackComment"] = props.pop("feedback_comment")

        # suggested_prompt → suggestedPrompt wire name
        if "suggested_prompt" in props:
            props["suggestedPrompt"] = props.pop("suggested_prompt")

        # reaction_type → wire name
        if "reaction_type" in props:
            props["reactionType"] = props.pop("reaction_type")

        # tool fields → wire names
        for old, new in [("tool_name", "toolName"), ("tool_input", "toolInput"),
                         ("tool_output", "toolOutput"), ("tool_status", "toolStatus"),
                         ("signal_name", "signalName"), ("signal_data", "signalData"),
                         ("message_span_id", "messageSpanId"),
                         ("prompt_tokens", "promptTokens"),
                         ("completion_tokens", "completionTokens"),
                         ("total_tokens", "totalTokens"),
                         ("cost_usd", "costUsd")]:
            if old in props:
                props[new] = props.pop(old)

        data["props"] = props
        return data


class EventBatch(BaseModel):
    """A batch of events sent in a single POST to ``/data/agentic``.

    The exporter collects finished spans, converts each to an AgentEvent,
    and ships them as one EventBatch.
    """

    events: List[AgentEvent] = Field(
        ...,
        description="List of agent events to ingest.",
        min_length=1,
    )

    def to_wire(self) -> Dict[str, Any]:
        """Serialize the batch for network export."""
        return {"events": [e.to_wire() for e in self.events]}
