"""
Session enrichment types for the Amplitude AI SDK.

Provides :class:`SessionEnrichments`, :class:`TopicClassification`,
:class:`RubricScore`, :class:`EvidenceQuote`, and :class:`MessageLabel` —
the structured data types used by :meth:`AmplitudeAI.track_session_end`
and :meth:`AmplitudeAI.track_session_enrichment`.

These types use the same vocabulary as Amplitude's enrichment taxonomy
framework (topic models + rubrics), ensuring property naming consistency
with server-side ``[Agent] Session Evaluation``, ``[Agent] Topic
Classification``, and ``[Agent] Score`` events.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any


# ======================================================================
# Building blocks
# ======================================================================


@dataclass
class MessageLabel:
    """A key-value label attached to a message event for filtering and segmentation.

    Labels are customer-defined — use whatever keys make sense for your product.

    Example::

        MessageLabel(key="flow", value="onboarding")
        MessageLabel(key="sentiment", value="negative", confidence=0.87)
    """

    key: str
    """Label key (e.g., ``"sentiment"``, ``"topic"``, ``"flow"``, ``"intent"``)."""

    value: str
    """Label value (e.g., ``"neutral"``, ``"onboarding"``)."""

    confidence: float | None = None
    """Optional confidence score between 0.0 and 1.0."""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict, omitting ``None`` fields."""
        result: dict[str, Any] = {"key": self.key, "value": self.value}
        if self.confidence is not None:
            result["confidence"] = self.confidence
        return result


@dataclass
class EvidenceQuote:
    """A quoted excerpt from the conversation supporting a rubric score.

    Example::

        EvidenceQuote(quote="47 page hits (success:true)", turn_index=17, role="tool")
    """

    quote: str
    """The quoted text from the conversation."""

    turn_index: int
    """0-based position in the conversation."""

    role: str | None = None
    """Role of the speaker (``"user"``, ``"assistant"``, ``"tool"``)."""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict, omitting ``None`` fields."""
        result: dict[str, Any] = {"quote": self.quote, "turn_index": self.turn_index}
        if self.role is not None:
            result["role"] = self.role
        return result


# ======================================================================
# Topic classification
# ======================================================================


@dataclass
class TopicClassification:
    """Result of classifying a session for a single topic model.

    Supports both single-select (MECE) and multi-select modes:

    - **Single-select**: set ``l1`` (e.g., ``TopicClassification(l1="diagnostic")``)
    - **Multi-select**: set ``values`` and ``primary``
      (e.g., ``TopicClassification(values=["charts", "cohorts"], primary="charts")``)
    - **Subcategories**: optionally set ``subcategories`` for finer-grained classification
    """

    l1: str | None = None
    """L1 value for single-select topic models (MECE)."""

    values: list[str] | None = None
    """L1 values for multi-select topic models."""

    primary: str | None = None
    """Primary value when using multi-select mode."""

    l2: str | None = None
    """.. deprecated:: Use ``subcategories`` instead.  Kept for backward compatibility."""

    subcategories: list[str] | None = None
    """Subcategory codes for finer classification within the primary topic."""

    topics_covered: list[str] | None = None
    """All topics discussed in the session (multi-topic sessions)."""

    outcomes_by_topic: dict[str, str] | None = None
    """Outcome per topic (e.g., ``{"page_views": "response_provided"}``)."""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict, omitting ``None`` fields.

        If ``subcategories`` is set, it is included.  If only the legacy
        ``l2`` scalar is set, it is wrapped in a list and emitted as
        ``subcategories`` for forward compatibility.  The raw ``l2`` key
        is still emitted when set so older consumers can read it.
        """
        result = {k: v for k, v in asdict(self).items() if v is not None}
        if "subcategories" not in result and "l2" in result:
            result["subcategories"] = [result["l2"]]
        return result


# ======================================================================
# Rubric scoring
# ======================================================================


@dataclass
class RubricScore:
    """Result of scoring a session on a single rubric.

    Example::

        RubricScore(name="task_completion", score=0.85)
        RubricScore(
            name="response_quality", score=0.92,
            rationale="Clear and accurate",
            evidence=[EvidenceQuote(quote="...", turn_index=5, role="assistant")],
        )
    """

    name: str
    """Rubric name (e.g., ``"task_completion"``, ``"response_quality"``)."""

    score: float
    """Score between 0.0 and 1.0."""

    rationale: str | None = None
    """Optional explanation for the score."""

    evidence: list[EvidenceQuote] | None = None
    """Quoted excerpts from the conversation supporting this score."""

    improvement_opportunities: str | None = None
    """Suggested improvements based on this evaluation."""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict, omitting ``None`` fields."""
        result: dict[str, Any] = {"name": self.name, "score": self.score}
        if self.rationale is not None:
            result["rationale"] = self.rationale
        if self.evidence is not None:
            result["evidence"] = [e.to_dict() for e in self.evidence]
        if self.improvement_opportunities is not None:
            result["improvement_opportunities"] = self.improvement_opportunities
        return result


# ======================================================================
# Session enrichments
# ======================================================================


@dataclass
class SessionEnrichments:
    """Customer-provided session classifications.

    Used with :meth:`AmplitudeAI.track_session_end` and
    :meth:`AmplitudeAI.track_session_enrichment` to attach structured
    metadata to an AI session.  Primarily used when
    ``content_mode="customer_enriched"`` — the customer provides their own
    classifications instead of relying on server-side enrichment.

    Example::

        enrichments = SessionEnrichments(
            topic_classifications={
                "query_intent": TopicClassification(l1="diagnostic"),
                "product_area": TopicClassification(
                    values=["charts", "cohorts"], primary="charts"
                ),
            },
            rubrics=[
                RubricScore(name="task_completion", score=0.85),
                RubricScore(name="response_quality", score=0.92),
            ],
            overall_outcome="response_provided",
            quality_score=0.85,
        )
    """

    # --- Topic Models ---
    topic_classifications: dict[str, TopicClassification] | None = None
    """Categorical classification per topic model."""

    # --- Rubrics ---
    rubrics: list[RubricScore] | None = None
    """Scored evaluation dimensions."""

    # --- Outcome ---
    overall_outcome: str | None = None
    """How the session ended (e.g., ``"response_provided"``, ``"abandoned"``)."""

    # --- Boolean flags ---
    has_task_failure: bool = False
    has_negative_feedback: bool = False
    has_data_quality_issues: bool = False
    has_technical_failure: bool = False

    # --- Supplementary data ---
    error_categories: list[str] | None = None
    behavioral_patterns: list[str] | None = None
    custom_metadata: dict[str, Any] | None = None
    """Arbitrary customer-defined metadata."""

    schema_version: str = "2.0"

    # --- Session-level scores (Amplitude understands natively) ---
    quality_score: float | None = None
    """Overall quality score (0.0–1.0).  Derived from response_quality rubric server-side."""

    sentiment_score: float | None = None
    """Overall sentiment score (0.0–1.0).  Derived from user_satisfaction rubric server-side."""

    # --- Failure detail ---
    task_failure_type: str | None = None
    """Type of failure (e.g., ``"access_denied"``, ``"unable_to_complete"``, ``"no_action_taken"``)."""

    task_failure_reason: str | None = None
    """Free-text explanation of the failure."""

    # --- Feedback and error detail ---
    negative_feedback_phrases: list[str] | None = None
    """Detected phrases indicating negative user feedback."""

    data_quality_issues: list[str] | None = None
    """Detected data quality issues (e.g., ``"no_data_returned"``, ``"property_not_found"``)."""

    technical_error_count: int | None = None
    """Number of technical errors during the session."""

    # --- Agent chain metadata ---
    agent_chain: list[str] | None = None
    """Ordered list of agents in the delegation chain."""

    root_agent_name: str | None = None
    """Entry-point agent that started the chain."""

    # --- Request classification ---
    request_complexity: str | None = None
    """Complexity of the user's request (``"simple"``, ``"moderate"``, ``"complex"``, ``"ambiguous"``)."""

    # --- Retrospective message labels ---
    message_labels: dict[str, list[MessageLabel]] | None = None
    """Labels attached to specific messages after the session.

    Keyed by ``message_id`` (the UUID returned by ``track_user_message()``
    or ``track_ai_message()``).  Used for post-hoc classifier output.
    """

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a flat dict suitable for Amplitude event properties.

        - ``None`` fields are omitted.
        - ``TopicClassification`` and ``RubricScore`` objects are converted
          to plain dicts.
        - Boolean flags are always included (even when *False*) since they
          are useful for filtering.
        """
        result: dict[str, Any] = {}

        if self.topic_classifications is not None:
            result["topic_classifications"] = {
                name: tc.to_dict()
                for name, tc in self.topic_classifications.items()
            }

        if self.rubrics is not None:
            result["rubrics"] = [r.to_dict() for r in self.rubrics]

        if self.overall_outcome is not None:
            result["overall_outcome"] = self.overall_outcome

        # Boolean flags — always include
        result["has_task_failure"] = self.has_task_failure
        result["has_negative_feedback"] = self.has_negative_feedback
        result["has_data_quality_issues"] = self.has_data_quality_issues
        result["has_technical_failure"] = self.has_technical_failure

        if self.error_categories is not None:
            result["error_categories"] = self.error_categories

        if self.behavioral_patterns is not None:
            result["behavioral_patterns"] = self.behavioral_patterns

        if self.custom_metadata is not None:
            result["custom_metadata"] = self.custom_metadata

        result["schema_version"] = self.schema_version

        # Session-level scores
        if self.quality_score is not None:
            result["quality_score"] = self.quality_score
        if self.sentiment_score is not None:
            result["sentiment_score"] = self.sentiment_score

        # Failure detail
        if self.task_failure_type is not None:
            result["task_failure_type"] = self.task_failure_type
        if self.task_failure_reason is not None:
            result["task_failure_reason"] = self.task_failure_reason

        # Feedback and error detail
        if self.negative_feedback_phrases is not None:
            result["negative_feedback_phrases"] = self.negative_feedback_phrases
        if self.data_quality_issues is not None:
            result["data_quality_issues"] = self.data_quality_issues
        if self.technical_error_count is not None:
            result["technical_error_count"] = self.technical_error_count

        # Agent chain
        if self.agent_chain is not None:
            result["agent_chain"] = self.agent_chain
        if self.root_agent_name is not None:
            result["root_agent_name"] = self.root_agent_name

        # Request classification
        if self.request_complexity is not None:
            result["request_complexity"] = self.request_complexity

        # Retrospective message labels
        if self.message_labels is not None:
            result["message_labels"] = {
                mid: [lbl.to_dict() for lbl in labels]
                for mid, labels in self.message_labels.items()
            }

        return result
