"""AgentEvent subclass that bypasses the base SDK's string truncation for $llm_message."""

from __future__ import annotations

import copy

from amplitude import BaseEvent


class AgentEvent(BaseEvent):
    """BaseEvent subclass that preserves ``$llm_message`` from SDK truncation.

    The base Amplitude Python SDK truncates all string values to 1024 chars
    in ``utils.truncate()``, called by ``get_event_body()``.  For [Agent]
    events, we need ``$llm_message`` content to pass through at full length
    (Nova already whitelists this property server-side).

    This subclass saves a deep copy of ``$llm_message`` before the parent's
    ``get_event_body()`` mutates it, then restores the original after.
    """

    def get_event_body(self) -> dict:
        ep = self.event_properties
        original_llm_message = (
            copy.deepcopy(ep["$llm_message"])
            if ep and "$llm_message" in ep
            else None
        )

        body = super().get_event_body()

        if original_llm_message is not None:
            if "event_properties" in body:
                body["event_properties"]["$llm_message"] = original_llm_message
            # Undo in-place mutation on self so repeated calls are safe
            if ep is not None:
                ep["$llm_message"] = original_llm_message

        return body
