"""
Memory-operation event handler mixin for NoveumCrewAIListener.

Handles CrewAI ``BaseEventListener`` memory events across all three memory
subsystems: short-term, long-term, and entity memory.

Each operation family follows the same started → completed/failed lifecycle:

  Query (read / search):
  - ``on_memory_query_started``     → open ``crewai.memory_op`` span;
                                       capture memory_type, query text, agent_role
  - ``on_memory_query_completed``   → close as SUCCESS; write results_count,
                                       duration_ms, full serialised results
  - ``on_memory_query_failed``      → close as ERROR

  Save (write / store):
  - ``on_memory_save_started``      → open ``crewai.memory_op`` span;
                                       capture memory_type, full value/metadata
  - ``on_memory_save_completed``    → close as SUCCESS; write duration_ms
  - ``on_memory_save_failed``       → close as ERROR

  Retrieval (bulk context pull — distinct from single-item query):
  - ``on_memory_retrieval_started``  → open ``crewai.memory_op`` span;
                                        capture task_id (the requesting task)
  - ``on_memory_retrieval_completed``→ close as SUCCESS; write full memory
                                        content and duration_ms
  - ``on_memory_retrieval_failed``   → close as ERROR

State consumed / mutated (declared in _CrewAIObserverState):
    _lock, _is_shutdown,
    _agent_spans, _task_spans, _memory_op_spans, _memory_op_start_times
"""

from __future__ import annotations

import logging
import traceback
from typing import Any, Optional

from noveum_trace.integrations.crewai.crewai_constants import (
    ATTR_AGENT_ROLE,
    ATTR_MEMORY_DURATION_MS,
    ATTR_MEMORY_OP_ID,
    ATTR_MEMORY_OPERATION,
    ATTR_MEMORY_QUERY,
    ATTR_MEMORY_RESULT_COUNT,
    ATTR_MEMORY_STATUS,
    ATTR_MEMORY_TYPE,
    ATTR_STATUS_ERROR,
    ATTR_STATUS_SUCCESS,
    SPAN_MEMORY_QUERY,
    SPAN_MEMORY_RETRIEVAL,
    SPAN_MEMORY_SAVE,
)
from noveum_trace.integrations.crewai.crewai_state import _CrewAIObserverMixinBase
from noveum_trace.integrations.crewai.crewai_utils import (
    finish_span_common,
    monotonic_now,
)
from noveum_trace.integrations.crewai.crewai_utils import (
    resolve_agent_id as _resolve_agent_id,
)
from noveum_trace.integrations.crewai.crewai_utils import (
    safe_getattr,
    safe_json_dumps,
)

logger = logging.getLogger(__name__)

# Memory operation type strings used for ATTR_MEMORY_OPERATION
_OP_QUERY = "query"
_OP_SAVE = "save"
_OP_RETRIEVAL = "retrieval"


class _MemoryHandlersMixin(_CrewAIObserverMixinBase):
    """
    Handler methods for CrewAI memory operation events.

    All public methods match the ``BaseEventListener`` callback signature::

        def on_memory_query_started(self, source, event): ...

    ``source`` is typically the memory object or the Agent; ``event`` carries
    the per-operation payload.  Every method is fully exception-shielded.

    All handlers no-op when ``capture_memory`` is ``False`` on the listener.
    """

    # =========================================================================
    # QUERY — started / completed / failed
    # =========================================================================

    def on_memory_query_started(self, source: Any, event: Any) -> None:
        """
        Open a ``crewai.memory_op`` span for a memory search/query operation.

        Attributes set at span open
        ---------------------------
        - ``memory.op_id``     — unique operation identifier
        - ``memory.type``      — ``"short_term"`` | ``"long_term"`` | ``"entity"``
                                 | ``"external"`` | ``"unknown"``
        - ``memory.operation`` — ``"query"``
        - ``memory.query``     — search query text (full)
        - ``agent.role``       — role of the querying agent (correlation)
        """
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            memory_type = _resolve_memory_type(source, event)
            agent_id = _resolve_agent_id(source, event)

            attrs: dict[str, Any] = {
                ATTR_MEMORY_OP_ID: op_id,
                ATTR_MEMORY_TYPE: memory_type,
                ATTR_MEMORY_OPERATION: _OP_QUERY,
            }

            query = (
                safe_getattr(event, "query")
                or safe_getattr(event, "search_query")
                or safe_getattr(event, "text")
            )
            if query:
                attrs[ATTR_MEMORY_QUERY] = str(query)

            agent_role = safe_getattr(event, "agent_role") or safe_getattr(
                source, "role"
            )
            if agent_role:
                attrs[ATTR_AGENT_ROLE] = str(agent_role)

            # memory.limit — max results requested (0 is valid; do not use ``or``)
            limit_raw = safe_getattr(event, "limit")
            limit = limit_raw if limit_raw is not None else safe_getattr(event, "top_k")
            if limit is not None:
                try:
                    attrs["memory.limit"] = int(limit)
                except (TypeError, ValueError):
                    pass

            # memory.score_threshold — minimum relevance score requested
            score_threshold = safe_getattr(event, "score_threshold")
            if score_threshold is None:
                score_threshold = safe_getattr(event, "threshold")
            if score_threshold is None:
                score_threshold = safe_getattr(event, "min_score")
            if score_threshold is not None:
                try:
                    attrs["memory.score_threshold"] = float(score_threshold)
                except (TypeError, ValueError):
                    pass

            task_id = _resolve_memory_task_id(source, event)
            crew_id = _resolve_memory_crew_id(source, event)
            self._open_memory_span(
                op_id,
                agent_id,
                attrs,
                span_name=SPAN_MEMORY_QUERY,
                task_id=task_id,
                crew_id=crew_id,
            )
            logger.debug(
                "Memory query span opened: op_id=%s type=%s", op_id, memory_type
            )
        except Exception:
            logger.debug("on_memory_query_started error:\n%s", traceback.format_exc())

    def on_memory_query_completed(self, source: Any, event: Any) -> None:
        """
        Close the memory query span as SUCCESS.

        Attributes written
        ------------------
        - ``memory.result_count``    — number of results returned
        - ``memory.results_preview`` — full JSON serialisation of results
        - ``memory.status``          — ``"success"``
        - ``memory.duration_ms``     — wall-clock duration
        """
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            extra: dict[str, Any] = {}

            results_raw = safe_getattr(event, "results")
            results = (
                results_raw
                if results_raw is not None
                else safe_getattr(event, "memories")
            )
            if results is not None:
                try:
                    count = len(results) if hasattr(results, "__len__") else None
                    if count is not None:
                        extra[ATTR_MEMORY_RESULT_COUNT] = count
                    preview = safe_json_dumps(results)
                    if preview:
                        extra["memory.results_preview"] = preview
                except Exception:
                    pass

            self._finish_memory_span(op_id, ATTR_STATUS_SUCCESS, None, extra)
        except Exception:
            logger.debug("on_memory_query_completed error:\n%s", traceback.format_exc())

    def on_memory_query_failed(self, source: Any, event: Any) -> None:
        """Close the memory query span as ERROR."""
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            error = safe_getattr(event, "error") or safe_getattr(event, "exception")
            self._finish_memory_span(op_id, ATTR_STATUS_ERROR, error)
        except Exception:
            logger.debug("on_memory_query_failed error:\n%s", traceback.format_exc())

    # =========================================================================
    # SAVE — started / completed / failed
    # =========================================================================

    def on_memory_save_started(self, source: Any, event: Any) -> None:
        """
        Open a ``crewai.memory_op`` span for a memory write/save operation.

        Attributes set at span open
        ---------------------------
        - ``memory.op_id``        — unique operation identifier
        - ``memory.type``         — memory subsystem type
        - ``memory.operation``    — ``"save"``
        - ``memory.value``        — the value being stored (full)
        - ``memory.metadata``     — full JSON of associated metadata when available
        - ``agent.role``          — role of the saving agent (correlation)
        """
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            memory_type = _resolve_memory_type(source, event)
            agent_id = _resolve_agent_id(source, event)

            attrs: dict[str, Any] = {
                ATTR_MEMORY_OP_ID: op_id,
                ATTR_MEMORY_TYPE: memory_type,
                ATTR_MEMORY_OPERATION: _OP_SAVE,
            }

            value = safe_getattr(event, "value")
            if value is None:
                value = safe_getattr(event, "content")
            if value is None:
                value = safe_getattr(event, "memory")
            if value is not None:
                if isinstance(value, str):
                    attrs["memory.value"] = value
                else:
                    attrs["memory.value"] = safe_json_dumps(value)

            metadata = safe_getattr(event, "metadata")
            if metadata is None:
                metadata = safe_getattr(event, "meta")
            if metadata is not None:
                attrs["memory.metadata"] = safe_json_dumps(metadata)

            agent_role = safe_getattr(event, "agent_role") or safe_getattr(
                source, "role"
            )
            if agent_role:
                attrs[ATTR_AGENT_ROLE] = str(agent_role)

            task_id = _resolve_memory_task_id(source, event)
            crew_id = _resolve_memory_crew_id(source, event)
            self._open_memory_span(
                op_id,
                agent_id,
                attrs,
                span_name=SPAN_MEMORY_SAVE,
                task_id=task_id,
                crew_id=crew_id,
            )
            logger.debug(
                "Memory save span opened: op_id=%s type=%s", op_id, memory_type
            )
        except Exception:
            logger.debug("on_memory_save_started error:\n%s", traceback.format_exc())

    def on_memory_save_completed(self, source: Any, event: Any) -> None:
        """Close the memory save span as SUCCESS."""
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            self._finish_memory_span(op_id, ATTR_STATUS_SUCCESS, None)
        except Exception:
            logger.debug("on_memory_save_completed error:\n%s", traceback.format_exc())

    def on_memory_save_failed(self, source: Any, event: Any) -> None:
        """Close the memory save span as ERROR."""
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            error = safe_getattr(event, "error") or safe_getattr(event, "exception")
            self._finish_memory_span(op_id, ATTR_STATUS_ERROR, error)
        except Exception:
            logger.debug("on_memory_save_failed error:\n%s", traceback.format_exc())

    # =========================================================================
    # RETRIEVAL — started / completed / failed
    # =========================================================================

    def on_memory_retrieval_started(self, source: Any, event: Any) -> None:
        """
        Open a ``crewai.memory_op`` span for a bulk context retrieval.

        Retrieval is distinct from query: CrewAI fires this when it pulls the
        full relevant memory context for a task (not a targeted search).  It
        may span multiple memory subsystems internally.

        Attributes set at span open
        ---------------------------
        - ``memory.op_id``        — unique operation identifier
        - ``memory.type``         — memory subsystem or ``"all"``
        - ``memory.operation``    — ``"retrieval"``
        - ``task.id``             — id of the requesting task (correlation)
        - ``agent.role``          — role of the agent context being hydrated
        """
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            memory_type = _resolve_memory_type(source, event, default="all")
            agent_id = _resolve_agent_id(source, event)

            attrs: dict[str, Any] = {
                ATTR_MEMORY_OP_ID: op_id,
                ATTR_MEMORY_TYPE: memory_type,
                ATTR_MEMORY_OPERATION: _OP_RETRIEVAL,
            }

            task_id = _resolve_memory_task_id(source, event)
            if task_id:
                attrs["task.id"] = task_id
                attrs["memory.task_id"] = task_id  # spec key

            agent_role = safe_getattr(event, "agent_role") or safe_getattr(
                source, "role"
            )
            if agent_role:
                attrs[ATTR_AGENT_ROLE] = str(agent_role)

            crew_id = _resolve_memory_crew_id(source, event)
            self._open_memory_span(
                op_id,
                agent_id,
                attrs,
                span_name=SPAN_MEMORY_RETRIEVAL,
                task_id=task_id,
                crew_id=crew_id,
            )
            logger.debug(
                "Memory retrieval span opened: op_id=%s type=%s", op_id, memory_type
            )
        except Exception:
            logger.debug(
                "on_memory_retrieval_started error:\n%s", traceback.format_exc()
            )

    def on_memory_retrieval_completed(self, source: Any, event: Any) -> None:
        """
        Close the memory retrieval span as SUCCESS.

        Attributes written
        ------------------
        - ``memory.content_preview`` — full retrieved memory context
        - ``memory.result_count``    — number of memory items retrieved (when countable)
        - ``memory.status``          — ``"success"``
        - ``memory.duration_ms``     — wall-clock duration
        """
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            extra: dict[str, Any] = {}

            memory_content = None
            for _attr in (
                "memory_content",
                "context",
                "memories",
                "results",
            ):
                val = safe_getattr(event, _attr)
                if val is not None:
                    memory_content = val
                    break
            if memory_content is not None:
                if isinstance(memory_content, str):
                    extra["memory.content_preview"] = memory_content
                else:
                    try:
                        count = (
                            len(memory_content)
                            if hasattr(memory_content, "__len__")
                            else None
                        )
                        if count is not None:
                            extra[ATTR_MEMORY_RESULT_COUNT] = count
                    except Exception:
                        pass
                    extra["memory.content_preview"] = safe_json_dumps(memory_content)

            self._finish_memory_span(op_id, ATTR_STATUS_SUCCESS, None, extra)
        except Exception:
            logger.debug(
                "on_memory_retrieval_completed error:\n%s", traceback.format_exc()
            )

    def on_memory_retrieval_failed(self, source: Any, event: Any) -> None:
        """Close the memory retrieval span as ERROR."""
        if not self._is_active() or not self.capture_memory:
            return
        try:
            op_id = _resolve_op_id(event, source)
            error = safe_getattr(event, "error") or safe_getattr(event, "exception")
            self._finish_memory_span(op_id, ATTR_STATUS_ERROR, error)
        except Exception:
            logger.debug(
                "on_memory_retrieval_failed error:\n%s", traceback.format_exc()
            )

    # =========================================================================
    # Internal helpers
    # =========================================================================

    def _open_memory_span(
        self,
        op_id: str,
        agent_id: Optional[str],
        attrs: dict[str, Any],
        span_name: str = SPAN_MEMORY_QUERY,
        *,
        task_id: Optional[str] = None,
        crew_id: Optional[str] = None,
    ) -> None:
        """Create a memory span with the given *span_name* and register it in state."""
        start_t = monotonic_now()

        # Parent: same resolution as LLM/A2A — agent span, else explicit task span
        # (never iterate arbitrary open tasks — avoids wrong trace under concurrency).
        parent_span = self._get_agent_or_task_span(agent_id, task_id)

        span = self._create_child_span(
            span_name,
            parent_span=parent_span,
            attributes=attrs,
            crew_id=crew_id,
            task_id=task_id,
        )

        with self._lock:
            self._memory_op_spans[op_id] = span
            self._memory_op_start_times[op_id] = start_t

    def _finish_memory_span(
        self,
        op_id: str,
        status: str,
        error: Any,
        extra_attrs: Optional[dict[str, Any]] = None,
    ) -> None:
        """Write final attributes onto the memory op span and close it."""
        with self._lock:
            span = self._memory_op_spans.pop(op_id, None)
            start_t = self._memory_op_start_times.pop(op_id, None)

        if span is None:
            logger.debug("_finish_memory_span: no open span for op_id=%s", op_id)
            return

        finish_span_common(
            span,
            start_t=start_t,
            status=status,
            status_attr=ATTR_MEMORY_STATUS,
            duration_attr=ATTR_MEMORY_DURATION_MS,
            error=error,
            extra_attrs=extra_attrs,
            log_label="_finish_memory_span",
        )

        logger.debug("Memory op span closed: op_id=%s status=%s", op_id, status)


# =============================================================================
# Module-level helpers (pure functions — no state access)
# =============================================================================


def _resolve_op_id(event: Any, source: Any) -> str:
    """Return a stable string key for pairing memory started ↔ completed/failed.

    CrewAI's event bus assigns ``started_event_id`` on scope-ending events to the
    ``event_id`` of the matching ``*_started`` event (see ``event_bus._prepare_event``).

    Falling back to Python ``id(event)`` **breaks** pairing because ``started`` and
    ``completed`` are different event instances: the completed handler never finds the
    open span, so the span stays in ``Trace.active_spans`` until ``Trace.finish()``,
    which force-finishes it with the **trace** end time — producing multi‑minute
    ``memory.*`` durations that track the whole kickoff instead of the real op.
    """
    started_link = safe_getattr(event, "started_event_id")
    if isinstance(started_link, str) and started_link.strip():
        return started_link
    return str(
        safe_getattr(event, "event_id")
        or safe_getattr(event, "memory_op_id")
        or safe_getattr(event, "op_id")
        or safe_getattr(event, "run_id")
        or id(event)
    )


def _resolve_memory_task_id(source: Any, event: Any) -> Optional[str]:
    """Return the task_id for this memory event, or ``None``."""
    raw = (
        safe_getattr(event, "task_id")
        or safe_getattr(safe_getattr(event, "task"), "id")
        or safe_getattr(safe_getattr(event, "from_task"), "id")
        or safe_getattr(source, "task_id")
        or safe_getattr(safe_getattr(source, "task"), "id")
    )
    return str(raw) if raw is not None else None


def _resolve_memory_crew_id(source: Any, event: Any) -> Optional[str]:
    """Return the crew_id for this memory event, or ``None``."""
    raw = (
        safe_getattr(event, "crew_id")
        or safe_getattr(source, "crew_id")
        or safe_getattr(safe_getattr(source, "crew"), "id")
        or safe_getattr(safe_getattr(event, "crew"), "id")
    )
    return str(raw) if raw is not None else None


def _resolve_memory_type(source: Any, event: Any, default: str = "unknown") -> str:
    """
    Resolve the memory subsystem type from event, source class name, or default.

    Maps raw strings to canonical values:
      ``"short_term"``  | ``"long_term"`` | ``"entity"`` |
      ``"external"``    | ``"all"``       | ``"unknown"``
    """
    raw = (
        safe_getattr(event, "memory_type")
        or safe_getattr(event, "type")
        or safe_getattr(source, "memory_type")
        or safe_getattr(source, "type")
    )
    if raw is None:
        # Infer from class name of source
        class_name = type(source).__name__.lower()
        if "short" in class_name:
            return "short_term"
        if "long" in class_name:
            return "long_term"
        if "entity" in class_name:
            return "entity"
        if "external" in class_name:
            return "external"
        return default

    raw_str = str(raw).lower().replace("-", "_").replace(" ", "_")
    for canonical in ("short_term", "long_term", "entity", "external", "all"):
        if canonical in raw_str:
            return canonical
    return raw_str or default