"""
Testing utilities for the Amplitude AI SDK.

Provides :class:`MockAmplitudeAI` — a drop-in replacement for
:class:`AmplitudeAI` that captures events in-memory instead of sending
them over the network.  Designed for use in unit and integration tests.

Example::

    from amplitude_ai.testing import MockAmplitudeAI

    mock = MockAmplitudeAI()
    mock.track_user_message(user_id="u1", content="Hi", session_id="s1")

    assert len(mock.events) == 1
    mock.assert_event_tracked("[Agent] User Message", user_id="u1")
"""

from __future__ import annotations

import time
from typing import Any
from unittest.mock import MagicMock

from amplitude import BaseEvent

from .client import AmplitudeAI
from .config import AIConfig
from .core.tracking import (
    EVENT_SESSION_END,
    PROP_AGENT_ID,
    PROP_SESSION_ID,
)


class MockAmplitudeAI(AmplitudeAI):
    """In-memory :class:`AmplitudeAI` for testing.

    All events are captured in :attr:`events` rather than sent to
    Amplitude's servers.

    Parameters
    ----------
    config:
        Optional :class:`AIConfig`.  Defaults to ``AIConfig()`` (full
        content mode, no PII redaction).
    """

    def __init__(self, config: AIConfig | None = None) -> None:
        # Create a real Amplitude instance with a dummy key so internal
        # code works, then redirect the underlying track to capture events.
        super().__init__(api_key="mock-api-key", config=config)

        self.events: list[BaseEvent] = []
        """All events tracked through this mock client."""

        # Replace the underlying Amplitude client's track method so the
        # _TrackingProxy (which calls self._amplitude.track) captures
        # events instead of sending them.  Debug/dry-run are already
        # handled by the proxy — no separate hook needed.
        self._amplitude.track = self._capture_event
        # Neutralise flush/shutdown so tests don't spawn threads
        self._amplitude.flush = MagicMock(return_value=[])
        self._amplitude.shutdown = MagicMock()

    # ------------------------------------------------------------------
    # Internal capture
    # ------------------------------------------------------------------

    def _capture_event(self, event: BaseEvent) -> None:
        """Intercept ``amplitude.track()`` and store the event."""
        self.events.append(event)

    # ------------------------------------------------------------------
    # Query helpers
    # ------------------------------------------------------------------

    def get_events(self, event_type: str | None = None) -> list[BaseEvent]:
        """Return captured events, optionally filtered by event type.

        Args:
            event_type: If provided, only return events matching this type
                        (e.g., ``"[Agent] User Message"``).
        """
        if event_type is None:
            return list(self.events)
        return [e for e in self.events if e.event_type == event_type]

    def assert_event_tracked(
        self,
        event_type: str,
        *,
        user_id: str | None = None,
        **expected_props: Any,
    ) -> BaseEvent:
        """Assert that at least one event of the given type was tracked.

        Optionally check ``user_id`` and event property values.  Returns
        the first matching event.

        Raises
        ------
        AssertionError
            If no matching event is found.
        """
        candidates = self.get_events(event_type)
        if not candidates:
            tracked_types = {e.event_type for e in self.events}
            raise AssertionError(
                f"No '{event_type}' event tracked. "
                f"Tracked types: {tracked_types or '(none)'}"
            )

        for event in candidates:
            if user_id is not None and event.user_id != user_id:
                continue

            if expected_props:
                props = event.event_properties or {}
                if all(props.get(k) == v for k, v in expected_props.items()):
                    return event
            else:
                return event

        # Build a helpful error message
        filters = []
        if user_id is not None:
            filters.append(f"user_id={user_id!r}")
        for k, v in expected_props.items():
            filters.append(f"{k}={v!r}")
        filter_desc = ", ".join(filters)
        raise AssertionError(
            f"Found {len(candidates)} '{event_type}' event(s) but none "
            f"matched filters: {filter_desc}"
        )

    # ------------------------------------------------------------------
    # Error / latency simulation
    # ------------------------------------------------------------------

    def simulate_error(self, error_type: str, message: str | None = None) -> None:
        """Configure the mock to raise on the next tracking call.

        Useful for testing error-handling paths.

        Args:
            error_type: Error category (e.g. ``"rate_limit"``, ``"network"``).
            message: Optional custom error message.
        """
        self._simulated_error: tuple[str, str] | None = (
            error_type,
            message or f"Simulated {error_type} error",
        )

        def _raise_on_track(event: BaseEvent) -> None:
            err = self._simulated_error
            if err:
                raise RuntimeError(f"[MockAmplitudeAI] {err[0]}: {err[1]}")

        self._amplitude.track = _raise_on_track  # type: ignore[assignment]

    def simulate_latency(self, seconds: float) -> None:
        """Add artificial latency (in seconds) to each tracking call.

        Useful for testing timeout and retry logic.
        """
        self._simulated_latency = seconds

        def _delayed_track(event: BaseEvent) -> None:
            if self._simulated_latency > 0:
                time.sleep(self._simulated_latency)
            self.events.append(event)

        self._amplitude.track = _delayed_track  # type: ignore[assignment]

    def clear_simulations(self) -> None:
        """Remove any simulated error or latency, restoring normal capture."""
        self._simulated_error = None
        self._simulated_latency = 0.0
        self._amplitude.track = self._capture_event  # type: ignore[assignment]

    def reset(self) -> None:
        """Clear all captured events, simulations, and turn counters."""
        self.events.clear()
        self._simulated_error = None
        self._simulated_latency = 0.0
        # Restore base capture.  Debug/dry-run is handled by the proxy.
        self._amplitude.track = self._capture_event  # type: ignore[assignment]
        with self._turn_lock:
            self._session_turn_counters.clear()

    # ------------------------------------------------------------------
    # Session / agent query helpers
    # ------------------------------------------------------------------

    def events_for_session(self, session_id: str) -> list[BaseEvent]:
        """Return all events for a specific session.

        Filters on the ``[Agent] Session ID`` event property.
        """
        return [
            e for e in self.events
            if (e.event_properties or {}).get(PROP_SESSION_ID) == session_id
        ]

    def events_for_agent(self, agent_id: str) -> list[BaseEvent]:
        """Return all events for a specific agent.

        Filters on the ``[Agent] Agent ID`` event property.
        """
        return [
            e for e in self.events
            if (e.event_properties or {}).get(PROP_AGENT_ID) == agent_id
        ]

    def assert_session_closed(self, session_id: str) -> BaseEvent:
        """Assert a ``[Agent] Session End`` event was tracked for *session_id*.

        Returns the first matching event.

        Raises
        ------
        AssertionError
            If no matching session-end event is found.
        """
        candidates = [
            e for e in self.events
            if e.event_type == EVENT_SESSION_END
            and (e.event_properties or {}).get(PROP_SESSION_ID) == session_id
        ]
        if not candidates:
            ended = [
                (e.event_properties or {}).get(PROP_SESSION_ID)
                for e in self.events
                if e.event_type == EVENT_SESSION_END
            ]
            raise AssertionError(
                f"No '{EVENT_SESSION_END}' event for session '{session_id}'. "
                f"Ended sessions: {ended or '(none)'}"
            )
        return candidates[0]
