"""
Tracking functions that create AgentEvent objects for Amplitude Agent analytics.
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any
import uuid

from amplitude import Amplitude

from .agent_event import AgentEvent
from .privacy import PrivacyConfig, _sanitize_structured_content, get_text_from_llm_message
from ..utils.logger import get_logger

if TYPE_CHECKING:
    from .enrichments import MessageLabel, SessionEnrichments

# Maximum length for serialized tool input/output to prevent overly large event properties
MAX_SERIALIZED_LENGTH = 10000

# ---------------------------------------------------------------------------
# Event type constants
# ---------------------------------------------------------------------------
EVENT_USER_MESSAGE = "[Agent] User Message"
EVENT_AI_RESPONSE = "[Agent] AI Response"
EVENT_TOOL_CALL = "[Agent] Tool Call"
EVENT_EMBEDDING = "[Agent] Embedding"
EVENT_SPAN = "[Agent] Span"
EVENT_SESSION_END = "[Agent] Session End"
EVENT_SESSION_ENRICHMENT = "[Agent] Session Enrichment"
EVENT_SCORE = "[Agent] Score"

# ---------------------------------------------------------------------------
# Property name constants
# ---------------------------------------------------------------------------
PROP_SESSION_ID = "[Agent] Session ID"
PROP_TRACE_ID = "[Agent] Trace ID"
PROP_TURN_ID = "[Agent] Turn ID"
PROP_MESSAGE_ID = "[Agent] Message ID"
PROP_MODEL_NAME = "[Agent] Model Name"
PROP_PROVIDER = "[Agent] Provider"
PROP_LATENCY_MS = "[Agent] Latency Ms"
PROP_TTFB_MS = "[Agent] TTFB Ms"
PROP_INPUT_TOKENS = "[Agent] Input Tokens"
PROP_OUTPUT_TOKENS = "[Agent] Output Tokens"
PROP_TOTAL_TOKENS = "[Agent] Total Tokens"
PROP_REASONING_TOKENS = "[Agent] Reasoning Tokens"
PROP_CACHE_READ_TOKENS = "[Agent] Cache Read Tokens"
PROP_CACHE_CREATION_TOKENS = "[Agent] Cache Creation Tokens"
PROP_COST_USD = "[Agent] Cost USD"
PROP_IS_ERROR = "[Agent] Is Error"
PROP_ERROR_MESSAGE = "[Agent] Error Message"
PROP_FINISH_REASON = "[Agent] Finish Reason"
PROP_TOOL_CALLS = "[Agent] Tool Calls"
PROP_TOOL_NAME = "[Agent] Tool Name"
PROP_TOOL_SUCCESS = "[Agent] Tool Success"
PROP_INVOCATION_ID = "[Agent] Invocation ID"
PROP_PARENT_MESSAGE_ID = "[Agent] Parent Message ID"
PROP_AGENT_ID = "[Agent] Agent ID"
PROP_PARENT_AGENT_ID = "[Agent] Parent Agent ID"
PROP_TOOL_INPUT = "[Agent] Tool Input"
PROP_TOOL_OUTPUT = "[Agent] Tool Output"
PROP_SDK_VERSION = "[Agent] SDK Version"
PROP_RUNTIME = "[Agent] Runtime"
PROP_ENV = "[Agent] Env"
PROP_LOCALE = "[Agent] Locale"
PROP_SPAN_KIND = "[Agent] Span Kind"
PROP_COMPONENT_TYPE = "[Agent] Component Type"
PROP_CUSTOMER_ORG_ID = "[Agent] Customer Org ID"
PROP_AGENT_VERSION = "[Agent] Agent Version"
PROP_AGENT_DESCRIPTION = "[Agent] Agent Description"
PROP_CONTEXT = "[Agent] Context"
PROP_MESSAGE_SOURCE = "[Agent] Message Source"

# New property constants (v0.3.0)
PROP_SPAN_ID = "[Agent] Span ID"
PROP_SPAN_NAME = "[Agent] Span Name"
PROP_PARENT_SPAN_ID = "[Agent] Parent Span ID"
PROP_INPUT_STATE = "[Agent] Input State"
PROP_OUTPUT_STATE = "[Agent] Output State"
PROP_EMBEDDING_DIMENSIONS = "[Agent] Embedding Dimensions"
PROP_ENRICHMENTS = "[Agent] Enrichments"
PROP_SCORE_NAME = "[Agent] Score Name"
PROP_SCORE_VALUE = "[Agent] Score Value"
PROP_TARGET_ID = "[Agent] Target ID"
PROP_TARGET_TYPE = "[Agent] Target Type"
PROP_EVALUATION_SOURCE = "[Agent] Evaluation Source"
PROP_COMMENT = "[Agent] Comment"
PROP_HAS_REASONING = "[Agent] Has Reasoning"
PROP_REASONING_CONTENT = "[Agent] Reasoning Content"
PROP_IDLE_TIMEOUT_MINUTES = "[Agent] Session Idle Timeout Minutes"

# Model configuration constants (v0.4.0)
PROP_SYSTEM_PROMPT = "[Agent] System Prompt"
PROP_SYSTEM_PROMPT_LENGTH = "[Agent] System Prompt Length"
PROP_TEMPERATURE = "[Agent] Temperature"
PROP_MAX_OUTPUT_TOKENS = "[Agent] Max Output Tokens"
PROP_TOP_P = "[Agent] Top P"
PROP_IS_STREAMING = "[Agent] Is Streaming"
PROP_PROMPT_ID = "[Agent] Prompt ID"

# Implicit feedback constants (v0.5.0)
PROP_IS_REGENERATION = "[Agent] Is Regeneration"
PROP_IS_EDIT = "[Agent] Is Edit"
PROP_EDITED_MESSAGE_ID = "[Agent] Edited Message ID"
PROP_WAS_COPIED = "[Agent] Was Copied"
PROP_ABANDONMENT_TURN = "[Agent] Abandonment Turn"

# Attachment constants (v0.5.0)
PROP_WAS_CACHED = "[Agent] Was Cached"
PROP_MODEL_TIER = "[Agent] Model Tier"

PROP_HAS_ATTACHMENTS = "[Agent] Has Attachments"
PROP_ATTACHMENT_TYPES = "[Agent] Attachment Types"
PROP_ATTACHMENT_COUNT = "[Agent] Attachment Count"
PROP_TOTAL_ATTACHMENT_SIZE = "[Agent] Total Attachment Size Bytes"
PROP_ATTACHMENTS = "[Agent] Attachments"

# Message label constants (v0.6.0)
PROP_MESSAGE_LABELS = "[Agent] Message Labels"
PROP_MESSAGE_LABEL_MAP = "[Agent] Message Label Map"

# Tool definitions (v0.8.0) — request-side tool schemas sent to the LLM
PROP_TOOL_DEFINITIONS = "[Agent] Tool Definitions"
PROP_TOOL_DEFINITIONS_COUNT = "[Agent] Tool Definitions Count"
PROP_TOOL_DEFINITIONS_HASH = "[Agent] Tool Definitions Hash"

# Session replay linking (v0.7.0)
PROP_SESSION_REPLAY_ID = "[Amplitude] Session Replay ID"

# Session-level enrichment property constants (v0.6.0)
PROP_QUALITY_SCORE = "[Agent] Quality Score"
PROP_SENTIMENT_SCORE = "[Agent] Sentiment Score"
PROP_TASK_FAILURE_TYPE = "[Agent] Task Failure Type"
PROP_TASK_FAILURE_REASON = "[Agent] Task Failure Reason"
PROP_AGENT_CHAIN = "[Agent] Agent Chain"
PROP_ROOT_AGENT_NAME = "[Agent] Root Agent Name"
PROP_REQUEST_COMPLEXITY = "[Agent] Request Complexity"

def _get_sdk_version() -> str:
    """Read version from package metadata, falling back to pyproject.toml on disk."""
    try:
        from importlib.metadata import version
        return version("amplitude-ai")
    except Exception:
        pass
    try:
        from pathlib import Path
        # __file__ is amplitude_ai/core/tracking.py -> .parent.parent.parent = amplitude_ai/ (outer)
        pyproject = Path(__file__).resolve().parent.parent.parent / "pyproject.toml"
        for line in pyproject.read_text().splitlines():
            if line.strip().startswith("version"):
                return line.split("=", 1)[1].strip().strip("\"'")
    except Exception:
        pass
    return "unknown"


SDK_VERSION = _get_sdk_version()


# ---------------------------------------------------------------------------
# Validation helpers (active only when PrivacyConfig.validate is True)
# ---------------------------------------------------------------------------

def _validate_required_str(value: Any, field_name: str) -> None:
    """Raise ValidationError if *value* is not a non-empty string."""
    from ..exceptions import ValidationError
    if not isinstance(value, str) or not value:
        raise ValidationError(f"{field_name} must be a non-empty string, got {value!r}")


def _validate_non_negative(value: Any, field_name: str) -> None:
    """Raise ValidationError if *value* is negative."""
    from ..exceptions import ValidationError
    if isinstance(value, (int, float)) and value < 0:
        raise ValidationError(f"{field_name} must be >= 0, got {value}")


def _validate_numeric(value: Any, field_name: str) -> None:
    """Raise ValidationError if *value* is not numeric."""
    from ..exceptions import ValidationError
    if not isinstance(value, (int, float)):
        raise ValidationError(f"{field_name} must be numeric, got {type(value).__name__}")


def _safe_size_bytes(attachment: dict) -> int:
    """Extract ``size_bytes`` from an attachment dict, coercing to ``int``.

    Returns ``0`` for missing, ``None``, or non-numeric values so that
    ``sum()`` never raises.
    """
    val = attachment.get("size_bytes", 0)
    if isinstance(val, (int, float)):
        return int(val)
    return 0


def _serialize_to_json_string(value: Any, max_length: int = MAX_SERIALIZED_LENGTH) -> str:
    """
    Serialize a value to a JSON string, truncating if necessary.

    This ensures complex objects (dicts, lists) are stored as a single string property
    rather than being flattened into multiple event properties by Amplitude.

    Args:
        value: The value to serialize
        max_length: Maximum length of the resulting string

    Returns:
        JSON string representation of the value, truncated if necessary
    """
    try:
        serialized = json.dumps(value, default=str, ensure_ascii=False)
    except (TypeError, ValueError, RecursionError):
        # Fallback to string representation if JSON serialization fails
        serialized = str(value)

    if len(serialized) > max_length:
        # Truncate and add indicator (ensure we don't create negative slice)
        truncate_at = max(0, max_length - 14)
        return serialized[:truncate_at] + "...[truncated]"
    return serialized


def _with_sdk_managed_properties(
    event_properties: dict[str, Any] | None,
    managed_properties: dict[str, Any],
) -> dict[str, Any]:
    """Merge event properties while ensuring SDK-managed keys take precedence."""
    return {**(event_properties or {}), **managed_properties}


def track_user_message(
    amplitude: Amplitude,
    user_id: str,
    message_content: str,
    session_id: str | None = None,
    trace_id: str | None = None,
    turn_id: int = 1,
    message_id: str | None = None,
    conversation_id: str | None = None,
    env: str | None = None,
    locale: str | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    # Implicit feedback signals
    is_regeneration: bool = False,
    is_edit: bool = False,
    edited_message_id: str | None = None,
    # Attachments (file uploads, images, etc.)
    attachments: list[dict[str, Any]] | None = None,
    # Inline message labels
    labels: list[MessageLabel] | None = None,
    # Message source (user vs agent delegation)
    message_source: str | None = None,
    event_properties: dict[str, Any] | None = None,
    user_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> str:
    """Track a user message event.

    Returns the ``message_id`` (UUID string) for this event.

    Args:
        is_regeneration: Whether the user requested the AI regenerate a
            previous response.
        is_edit: Whether the user edited a previous message and resubmitted.
        edited_message_id: The ``message_id`` of the original message that
            was edited (links the edit to the original).
        attachments: List of attachment metadata dicts.  Each dict should
            contain ``type`` (e.g., ``"pdf"``, ``"image"``, ``"csv"``),
            ``name`` (filename), and optionally ``size_bytes``,
            ``mime_type``, ``page_count``, ``dimensions``.
            Only metadata is sent -- file content is never transmitted.
        labels: Optional list of :class:`MessageLabel` instances attached
            to this message.  Emitted as ``[Agent] Message Labels``.
        message_source: ``"user"`` for real user messages, ``"agent"``
            for inter-agent delegation messages.
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_required_str(session_id, "session_id")

    if not message_id:
        message_id = str(uuid.uuid4())

    content_data = privacy_config.sanitize_content(message_content)

    properties = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_TURN_ID: turn_id,
            PROP_MESSAGE_ID: message_id,
            PROP_COMPONENT_TYPE: "user_input",
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if session_id:
        properties[PROP_SESSION_ID] = session_id
    if trace_id:
        properties[PROP_TRACE_ID] = trace_id
    if conversation_id and not session_id:
        properties[PROP_SESSION_ID] = conversation_id
    if env:
        properties[PROP_ENV] = env
    if locale:
        properties[PROP_LOCALE] = locale
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)

    # Implicit feedback signals
    if is_regeneration:
        properties[PROP_IS_REGENERATION] = True
    if is_edit:
        properties[PROP_IS_EDIT] = True
    if edited_message_id:
        properties[PROP_EDITED_MESSAGE_ID] = edited_message_id

    # Attachment metadata
    if attachments:
        properties[PROP_HAS_ATTACHMENTS] = True
        properties[PROP_ATTACHMENT_COUNT] = len(attachments)
        attachment_types = list({a.get("type", "unknown") for a in attachments})
        properties[PROP_ATTACHMENT_TYPES] = attachment_types
        total_size = sum(_safe_size_bytes(a) for a in attachments)
        if total_size > 0:
            properties[PROP_TOTAL_ATTACHMENT_SIZE] = total_size
        # Store full attachment metadata as serialized JSON
        properties[PROP_ATTACHMENTS] = _serialize_to_json_string(attachments)

    # Inline message labels
    if labels:
        properties[PROP_MESSAGE_LABELS] = _serialize_to_json_string(
            [lbl.to_dict() for lbl in labels]
        )

    if message_source:
        properties[PROP_MESSAGE_SOURCE] = message_source

    properties.update(content_data)

    event = AgentEvent(
        event_type=EVENT_USER_MESSAGE,
        user_id=user_id,
        event_properties=properties,
        user_properties=user_properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_USER_MESSAGE} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_USER_MESSAGE}: {e}")

    return message_id


def track_ai_message(
    amplitude: Amplitude,
    user_id: str,
    model_name: str,
    provider: str,
    response_content: str,
    latency_ms: float,
    session_id: str | None = None,
    trace_id: str | None = None,
    turn_id: int = 2,
    message_id: str | None = None,
    conversation_id: str | None = None,
    input_tokens: int | None = None,
    output_tokens: int | None = None,
    total_tokens: int | None = None,
    reasoning_tokens: int | None = None,
    cache_read_input_tokens: int | None = None,
    cache_creation_input_tokens: int | None = None,
    total_cost_usd: float | None = None,
    provider_ttfb_ms: float | None = None,
    is_error: bool = False,
    error_message: str | None = None,
    finish_reason: str | None = None,
    tool_calls: list[dict[str, Any]] | None = None,
    reasoning_content: str | None = None,
    # Tool definitions (request-side tool schemas sent to the LLM)
    tool_definitions: list[dict[str, Any]] | None = None,
    # Model configuration
    system_prompt: str | None = None,
    temperature: float | None = None,
    max_output_tokens: int | None = None,
    top_p: float | None = None,
    is_streaming: bool | None = None,
    prompt_id: str | None = None,
    # Implicit feedback signals
    was_copied: bool = False,
    # Semantic cache hit (full-response cache, distinct from token-level prompt caching)
    was_cached: bool = False,
    # Model tier override (auto-inferred from model_name if not provided)
    model_tier: str | None = None,
    # Attachments (AI-generated images, charts, files, etc.)
    attachments: list[dict[str, Any]] | None = None,
    # Inline message labels
    labels: list[MessageLabel] | None = None,
    env: str | None = None,
    locale: str | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    span_kind: str | None = None,
    event_properties: dict[str, Any] | None = None,
    user_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> str:
    """Track an AI response event.

    Returns the ``message_id`` (UUID string) for this event.

    Args:
        was_copied: Whether the user copied the AI response content.
        attachments: List of attachment dicts for AI-generated rich media.
            Include a ``url`` key to enable rendering in the LLM session
            viewer.  Only metadata is sent -- file content is never
            transmitted.
        labels: Optional list of :class:`MessageLabel` instances attached
            to this message.  Emitted as ``[Agent] Message Labels``.
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_required_str(model_name, "model")
        _validate_non_negative(latency_ms, "latency_ms")

    if not message_id:
        message_id = str(uuid.uuid4())

    content_data = privacy_config.sanitize_content(response_content)

    properties = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_TURN_ID: turn_id,
            PROP_MESSAGE_ID: message_id,
            PROP_MODEL_NAME: model_name,
            PROP_PROVIDER: provider,
            PROP_LATENCY_MS: latency_ms,
            PROP_IS_ERROR: is_error,
            PROP_COMPONENT_TYPE: "llm",
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if session_id:
        properties[PROP_SESSION_ID] = session_id
    if trace_id:
        properties[PROP_TRACE_ID] = trace_id
    if conversation_id and not session_id:
        properties[PROP_SESSION_ID] = conversation_id
    if env:
        properties[PROP_ENV] = env
    if locale:
        properties[PROP_LOCALE] = locale
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if span_kind:
        properties[PROP_SPAN_KIND] = span_kind

    # Implicit feedback signals
    if was_copied:
        properties[PROP_WAS_COPIED] = True
    if was_cached:
        properties[PROP_WAS_CACHED] = True

    # Model tier (user override or auto-inferred)
    if model_tier is not None:
        properties[PROP_MODEL_TIER] = model_tier
    else:
        from ..utils.model_tiers import infer_model_tier
        properties[PROP_MODEL_TIER] = infer_model_tier(model_name)

    properties.update(content_data)

    # Token usage
    if input_tokens is not None:
        properties[PROP_INPUT_TOKENS] = input_tokens
    if output_tokens is not None:
        properties[PROP_OUTPUT_TOKENS] = output_tokens
    if total_tokens is not None:
        properties[PROP_TOTAL_TOKENS] = total_tokens
    if reasoning_tokens is not None:
        properties[PROP_REASONING_TOKENS] = reasoning_tokens
    if cache_read_input_tokens is not None:
        properties[PROP_CACHE_READ_TOKENS] = cache_read_input_tokens
    if cache_creation_input_tokens is not None:
        properties[PROP_CACHE_CREATION_TOKENS] = cache_creation_input_tokens

    # Cost and performance
    if total_cost_usd is not None:
        properties[PROP_COST_USD] = total_cost_usd
    if provider_ttfb_ms is not None:
        properties[PROP_TTFB_MS] = provider_ttfb_ms
    if finish_reason is not None:
        properties[PROP_FINISH_REASON] = finish_reason
    if error_message is not None:
        properties[PROP_ERROR_MESSAGE] = error_message
    if tool_calls is not None:
        # Serialize tool_calls to JSON string to prevent property sprawl
        properties[PROP_TOOL_CALLS] = _serialize_to_json_string(tool_calls)

    # Reasoning/thinking content.  Pass reasoning_tokens so the method
    # can set PROP_HAS_REASONING when tokens are present even without
    # content (e.g., manual tracking where only token counts are available).
    # PROP_REASONING_TOKENS is already set above but the duplicate is harmless.
    reasoning_props = privacy_config.sanitize_reasoning_content(
        reasoning_content, reasoning_tokens=reasoning_tokens
    )
    properties.update(reasoning_props)

    # Model configuration
    system_prompt_props = privacy_config.sanitize_system_prompt(system_prompt)
    properties.update(system_prompt_props)

    # Tool definitions (request-side schemas)
    tool_def_props = privacy_config.sanitize_tool_definitions(tool_definitions)
    properties.update(tool_def_props)

    if temperature is not None:
        properties[PROP_TEMPERATURE] = temperature
    if max_output_tokens is not None:
        properties[PROP_MAX_OUTPUT_TOKENS] = max_output_tokens
    if top_p is not None:
        properties[PROP_TOP_P] = top_p
    if is_streaming is not None:
        properties[PROP_IS_STREAMING] = is_streaming
    if prompt_id is not None:
        properties[PROP_PROMPT_ID] = prompt_id

    # Attachment metadata
    if attachments:
        properties[PROP_HAS_ATTACHMENTS] = True
        properties[PROP_ATTACHMENT_COUNT] = len(attachments)
        attachment_types = list({a.get("type", "unknown") for a in attachments})
        properties[PROP_ATTACHMENT_TYPES] = attachment_types
        total_size = sum(_safe_size_bytes(a) for a in attachments)
        if total_size > 0:
            properties[PROP_TOTAL_ATTACHMENT_SIZE] = total_size
        properties[PROP_ATTACHMENTS] = _serialize_to_json_string(attachments)

    # Inline message labels
    if labels:
        properties[PROP_MESSAGE_LABELS] = _serialize_to_json_string(
            [lbl.to_dict() for lbl in labels]
        )

    event = AgentEvent(
        event_type=EVENT_AI_RESPONSE,
        user_id=user_id,
        event_properties=properties,
        user_properties=user_properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_AI_RESPONSE} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_AI_RESPONSE}: {e}")

    return message_id


def track_tool_call(
    amplitude: Amplitude,
    user_id: str,
    tool_name: str,
    success: bool,
    latency_ms: float,
    session_id: str | None = None,
    trace_id: str | None = None,
    turn_id: int = 1,
    tool_input: Any | None = None,
    tool_output: Any | None = None,
    invocation_id: str | None = None,
    conversation_id: str | None = None,
    parent_message_id: str | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    error_message: str | None = None,
    env: str | None = None,
    locale: str | None = None,
    span_kind: str | None = None,
    event_properties: dict[str, Any] | None = None,
    user_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> str:
    """Track a tool call event.

    Returns the ``invocation_id`` (UUID string) for this event.
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_required_str(tool_name, "tool_name")
        _validate_non_negative(latency_ms, "latency_ms")

    if not invocation_id:
        invocation_id = str(uuid.uuid4())

    properties = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_TURN_ID: turn_id,
            PROP_INVOCATION_ID: invocation_id,
            PROP_TOOL_NAME: tool_name,
            PROP_TOOL_SUCCESS: success,
            PROP_IS_ERROR: not success,
            PROP_LATENCY_MS: latency_ms,
            PROP_COMPONENT_TYPE: "tool",
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if session_id:
        properties[PROP_SESSION_ID] = session_id
    if trace_id:
        properties[PROP_TRACE_ID] = trace_id
    if conversation_id and not session_id:
        properties[PROP_SESSION_ID] = conversation_id
    if parent_message_id:
        properties[PROP_PARENT_MESSAGE_ID] = parent_message_id
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if error_message:
        properties[PROP_ERROR_MESSAGE] = error_message
    if env:
        properties[PROP_ENV] = env
    if locale:
        properties[PROP_LOCALE] = locale
    if span_kind:
        properties[PROP_SPAN_KIND] = span_kind

    # Handle tool input/output with privacy.
    # Like span state, tool data is arbitrary structured content — we apply
    # _sanitize_structured_content (recursive PII redaction on strings) rather
    # than sanitize_content() (which is designed for LLM messages).
    _effective_mode = privacy_config.content_mode
    if _effective_mode is None:
        _effective_mode = "metadata_only" if privacy_config.privacy_mode else "full"

    if tool_input is not None:
        if _effective_mode == "full":
            sanitized = _sanitize_structured_content(tool_input, redact_pii=privacy_config.redact_pii)
            properties[PROP_TOOL_INPUT] = _serialize_to_json_string(sanitized)
        # metadata_only / customer_enriched → omitted

    if tool_output is not None:
        if _effective_mode == "full":
            sanitized = _sanitize_structured_content(tool_output, redact_pii=privacy_config.redact_pii)
            properties[PROP_TOOL_OUTPUT] = _serialize_to_json_string(sanitized)
        # metadata_only / customer_enriched → omitted

    event = AgentEvent(
        event_type=EVENT_TOOL_CALL,
        user_id=user_id,
        event_properties=properties,
        user_properties=user_properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_TOOL_CALL} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_TOOL_CALL}: {e}")

    return invocation_id


def track_conversation(
    amplitude: Amplitude,
    user_id: str,
    messages: list[dict[str, Any]],
    session_id: str | None = None,
    conversation_id: str | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    event_properties: dict[str, Any] | None = None,
    user_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> None:
    """
    Track a complete conversation with alternating user and AI messages.

    Args:
        amplitude: Amplitude client instance
        user_id: Unique identifier for the user
        messages: List of message dictionaries with 'role' and 'content' fields
        session_id: Session identifier for grouping conversations
        conversation_id: Deprecated alias for session_id
        agent_id: Agent identifier
        parent_agent_id: Delegating agent (multi-agent orchestration)
        customer_org_id: End-customer org (multi-tenant platforms)
        agent_version: Agent code version (e.g., "2.1.0")
        context: Arbitrary segmentation context dict
        env: Environment (prod, staging, dev)
        event_properties: Additional event properties for all messages
        user_properties: User properties to set
        groups: Group analytics properties
    """
    effective_session_id = session_id or conversation_id or str(uuid.uuid4())

    turn_id = 1
    for message in messages:
        role = message.get("role", "")
        content = message.get("content", "")

        if role == "user":
            track_user_message(
                amplitude=amplitude,
                user_id=user_id,
                message_content=content,
                session_id=effective_session_id,
                turn_id=turn_id,
                agent_id=agent_id,
                parent_agent_id=parent_agent_id,
                customer_org_id=customer_org_id,
                agent_version=agent_version,
                description=description,
                context=context,
                env=env,
                event_properties=event_properties,
                user_properties=user_properties,
                groups=groups,
                privacy_config=privacy_config,
            )
        elif role in ["assistant", "ai"]:
            model_name = message.get("model", "unknown")
            provider = message.get("provider", "unknown")
            latency_ms = message.get("latency_ms", 0.0)

            track_ai_message(
                amplitude=amplitude,
                user_id=user_id,
                model_name=model_name,
                provider=provider,
                response_content=content,
                latency_ms=latency_ms,
                session_id=effective_session_id,
                turn_id=turn_id,
                input_tokens=message.get("input_tokens"),
                output_tokens=message.get("output_tokens"),
                total_tokens=message.get("total_tokens"),
                total_cost_usd=message.get("total_cost_usd"),
                agent_id=agent_id,
                parent_agent_id=parent_agent_id,
                customer_org_id=customer_org_id,
                agent_version=agent_version,
                description=description,
                context=context,
                env=env,
                event_properties=event_properties,
                user_properties=user_properties,
                groups=groups,
                privacy_config=privacy_config,
            )

        turn_id += 1


# ---------------------------------------------------------------------------
# New tracking functions (v0.3.0)
# ---------------------------------------------------------------------------


def track_embedding(
    amplitude: Amplitude,
    user_id: str,
    model: str,
    provider: str,
    latency_ms: float,
    input_tokens: int | None = None,
    dimensions: int | None = None,
    total_cost_usd: float | None = None,
    session_id: str | None = None,
    trace_id: str | None = None,
    turn_id: int | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    event_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,  # kept for API compat; unused
) -> str:
    """Track an embedding operation.

    Returns a ``span_id`` (UUID string) for this event.
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_non_negative(latency_ms, "latency_ms")

    span_id = str(uuid.uuid4())

    properties: dict[str, Any] = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_SPAN_ID: span_id,
            PROP_MODEL_NAME: model,
            PROP_PROVIDER: provider,
            PROP_LATENCY_MS: latency_ms,
            PROP_COMPONENT_TYPE: "embedding",
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if session_id:
        properties[PROP_SESSION_ID] = session_id
    if trace_id is not None:
        properties[PROP_TRACE_ID] = trace_id
    if turn_id is not None:
        properties[PROP_TURN_ID] = turn_id
    if input_tokens is not None:
        properties[PROP_INPUT_TOKENS] = input_tokens
    if dimensions is not None:
        properties[PROP_EMBEDDING_DIMENSIONS] = dimensions
    if total_cost_usd is not None:
        properties[PROP_COST_USD] = total_cost_usd
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if env:
        properties[PROP_ENV] = env

    event = AgentEvent(
        event_type=EVENT_EMBEDDING,
        user_id=user_id,
        event_properties=properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_EMBEDDING} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_EMBEDDING}: {e}")

    return span_id


def track_span(
    amplitude: Amplitude,
    user_id: str,
    span_name: str,
    trace_id: str,
    latency_ms: float,
    input_state: dict[str, Any] | None = None,
    output_state: dict[str, Any] | None = None,
    parent_span_id: str | None = None,
    is_error: bool = False,
    error_message: str | None = None,
    session_id: str | None = None,
    turn_id: int | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    event_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> str:
    """Track a generic operation span (e.g., vector search, rerank).

    Returns a ``span_id`` (UUID string) for this event.
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_non_negative(latency_ms, "latency_ms")

    span_id = str(uuid.uuid4())

    properties: dict[str, Any] = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_SPAN_ID: span_id,
            PROP_SPAN_NAME: span_name,
            PROP_TRACE_ID: trace_id,
            PROP_LATENCY_MS: latency_ms,
            PROP_IS_ERROR: is_error,
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if session_id:
        properties[PROP_SESSION_ID] = session_id
    if turn_id is not None:
        properties[PROP_TURN_ID] = turn_id
    if parent_span_id:
        properties[PROP_PARENT_SPAN_ID] = parent_span_id
    if error_message:
        properties[PROP_ERROR_MESSAGE] = error_message
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if env:
        properties[PROP_ENV] = env

    # Handle input_state/output_state with privacy — same approach as
    # track_tool_call.  Span state data is arbitrary structured JSON, not
    # LLM messages.  Routing through sanitize_content() would extract only
    # "text"/"content"/"message" keys (designed for LLM messages), silently
    # dropping other fields from the dict.
    _effective_mode = privacy_config.content_mode
    if _effective_mode is None:
        _effective_mode = "metadata_only" if privacy_config.privacy_mode else "full"

    if input_state is not None:
        if _effective_mode == "full":
            sanitized = _sanitize_structured_content(input_state, redact_pii=privacy_config.redact_pii)
            properties[PROP_INPUT_STATE] = _serialize_to_json_string(sanitized)
        # else: metadata_only/customer_enriched → omitted

    if output_state is not None:
        if _effective_mode == "full":
            sanitized = _sanitize_structured_content(output_state, redact_pii=privacy_config.redact_pii)
            properties[PROP_OUTPUT_STATE] = _serialize_to_json_string(sanitized)
        # else: metadata_only/customer_enriched → omitted

    event = AgentEvent(
        event_type=EVENT_SPAN,
        user_id=user_id,
        event_properties=properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_SPAN} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_SPAN}: {e}")

    return span_id


def track_session_end(
    amplitude: Amplitude,
    user_id: str,
    session_id: str,
    enrichments: SessionEnrichments | None = None,
    trace_id: str | None = None,
    turn_id: int | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    abandonment_turn: int | None = None,
    idle_timeout_minutes: int | None = None,
    event_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,  # kept for API compat; unused
) -> None:
    """Track an explicit session close, optionally with enrichments.

    Args:
        enrichments: Optional :class:`SessionEnrichments` instance.
        abandonment_turn: If set, indicates the user abandoned the session
            at this turn.  The value is the ``turn_id`` of the last user
            message that received an AI response before the user left.
            A session where ``abandonment_turn`` equals 1 (or is very low)
            strongly signals that the first AI response did not satisfy.
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_required_str(session_id, "session_id")

    properties: dict[str, Any] = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_SESSION_ID: session_id,
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if trace_id is not None:
        properties[PROP_TRACE_ID] = trace_id
    if turn_id is not None:
        properties[PROP_TURN_ID] = turn_id
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if env:
        properties[PROP_ENV] = env
    if abandonment_turn is not None:
        properties[PROP_ABANDONMENT_TURN] = abandonment_turn
    if idle_timeout_minutes is not None:
        properties[PROP_IDLE_TIMEOUT_MINUTES] = idle_timeout_minutes

    if enrichments is not None:
        enrichment_dict = enrichments.to_dict() if hasattr(enrichments, "to_dict") else enrichments
        properties[PROP_ENRICHMENTS] = _serialize_to_json_string(enrichment_dict)

    event = AgentEvent(
        event_type=EVENT_SESSION_END,
        user_id=user_id,
        event_properties=properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_SESSION_END} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_SESSION_END}: {e}")


def track_session_enrichment(
    amplitude: Amplitude,
    user_id: str,
    session_id: str,
    enrichments: SessionEnrichments,
    trace_id: str | None = None,
    turn_id: int | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    event_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,  # kept for API compat; unused
) -> None:
    """Track customer-provided session classifications.

    Used when ``content_mode="customer_enriched"`` — the customer provides
    their own topic classifications, rubric scores, and metadata instead
    of relying on server-side enrichment.

    Args:
        enrichments: A :class:`SessionEnrichments` instance (required).
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_required_str(session_id, "session_id")

    properties: dict[str, Any] = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_SESSION_ID: session_id,
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if trace_id is not None:
        properties[PROP_TRACE_ID] = trace_id
    if turn_id is not None:
        properties[PROP_TURN_ID] = turn_id
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if env:
        properties[PROP_ENV] = env

    enrichment_dict = enrichments.to_dict() if hasattr(enrichments, "to_dict") else enrichments
    properties[PROP_ENRICHMENTS] = _serialize_to_json_string(enrichment_dict)

    event = AgentEvent(
        event_type=EVENT_SESSION_ENRICHMENT,
        user_id=user_id,
        event_properties=properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_SESSION_ENRICHMENT} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_SESSION_ENRICHMENT}: {e}")


def track_score(
    amplitude: Amplitude,
    user_id: str,
    name: str,
    value: float,
    target_id: str,
    target_type: str = "message",
    source: str = "user",
    comment: str | None = None,
    session_id: str | None = None,
    trace_id: str | None = None,
    turn_id: int | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    description: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    event_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> None:
    """Track a score (user feedback, automated eval, or human annotation).

    Args:
        name: Score name (e.g., ``"user-feedback"``, ``"accuracy"``).
        value: Numeric value (binary 1/0, continuous 0-1, or rating 1-5).
        target_id: The ``message_id`` or ``session_id`` being scored.
        target_type: ``"message"`` or ``"session"``.
        source: ``"user"``, ``"ai"``, or ``"reviewer"``.
        comment: Optional text explanation (respects ``content_mode``).
    """
    privacy_config = privacy_config or PrivacyConfig()
    if privacy_config.validate:
        _validate_required_str(user_id, "user_id")
        _validate_numeric(value, "value")

    properties: dict[str, Any] = _with_sdk_managed_properties(
        event_properties,
        {
            PROP_SCORE_NAME: name,
            PROP_SCORE_VALUE: value,
            PROP_TARGET_ID: target_id,
            PROP_TARGET_TYPE: target_type,
            PROP_EVALUATION_SOURCE: source,
            PROP_SDK_VERSION: SDK_VERSION,
            PROP_RUNTIME: "python",
        },
    )

    if session_id:
        properties[PROP_SESSION_ID] = session_id
    if trace_id is not None:
        properties[PROP_TRACE_ID] = trace_id
    if turn_id is not None:
        properties[PROP_TURN_ID] = turn_id
    if agent_id:
        properties[PROP_AGENT_ID] = agent_id
    if parent_agent_id:
        properties[PROP_PARENT_AGENT_ID] = parent_agent_id
    if customer_org_id:
        properties[PROP_CUSTOMER_ORG_ID] = customer_org_id
    if agent_version:
        properties[PROP_AGENT_VERSION] = agent_version
    if description:
        properties[PROP_AGENT_DESCRIPTION] = description
    if context:
        properties[PROP_CONTEXT] = _serialize_to_json_string(context)
    if env:
        properties[PROP_ENV] = env

    # Comment respects content_mode
    if comment is not None:
        comment_data = privacy_config.sanitize_content(comment)
        if "$llm_message" in comment_data:
            properties[PROP_COMMENT] = get_text_from_llm_message(comment_data["$llm_message"])
        # metadata_only / customer_enriched → no comment

    event = AgentEvent(
        event_type=EVENT_SCORE,
        user_id=user_id,
        event_properties=properties,
        groups=groups,
    )

    try:
        amplitude.track(event)
        get_logger(amplitude).debug(f"Tracked {EVENT_SCORE} for user {user_id}")
    except Exception as e:
        get_logger(amplitude).error(f"Failed to track {EVENT_SCORE}: {e}")
