"""
CrewAI integration for Amplitude AI SDK.

Provides :class:`AmplitudeCrewAIHooks` — registers CrewAI before/after
hooks for LLM calls and tool calls to automatically track usage in
Amplitude.

Usage::

    from crewai import Crew, Agent, Task
    from amplitude import Amplitude
    from amplitude_ai.integrations.crewai import AmplitudeCrewAIHooks

    amplitude = Amplitude("YOUR_API_KEY")
    hooks = AmplitudeCrewAIHooks(
        amplitude=amplitude,
        user_id="user-1",
    )
    hooks.activate()

    crew = Crew(agents=[...], tasks=[...])
    result = crew.kickoff()

    hooks.deactivate()
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Dict, Optional

try:
    from crewai.utilities.events.event_listener import EventListener
    from crewai.utilities.events import crewai_event_bus

    CREWAI_AVAILABLE = True
except ImportError:
    CREWAI_AVAILABLE = False

from amplitude import Amplitude

from ..context import get_active_context
from ..core.privacy import PrivacyConfig
from ..core.tracking import (
    track_ai_message,
    track_tool_call,
    track_user_message,
)
from ..utils.costs import calculate_cost, infer_provider
from ..utils.logger import get_logger
from ..utils.providers import infer_provider_from_model


class AmplitudeCrewAIHooks:
    """CrewAI event listener that tracks LLM and tool usage in Amplitude.

    Registers callbacks on CrewAI's event bus to capture:

    - ``LLMCallStartedEvent`` / ``LLMCallCompletedEvent``
      -> ``[Agent] AI Response``
    - ``ToolUsageStartedEvent`` / ``ToolUsageFinishedEvent``
      -> ``[Agent] Tool Call``

    Token counts are extracted from CrewAI's ``LLMCallCompletedEvent``
    ``token_usage`` when available.

    Note: CrewAI events do not expose the original user prompt, so
    ``[Agent] User Message`` events are **not** emitted automatically.
    Call :func:`~amplitude_ai.core.tracking.track_user_message` manually
    before ``crew.kickoff()`` if you need to capture user input.
    """

    def __init__(
        self,
        amplitude: Amplitude,
        user_id: str,
        privacy_config: Optional[PrivacyConfig] = None,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        parent_agent_id: Optional[str] = None,
        customer_org_id: Optional[str] = None,
        agent_version: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
        env: Optional[str] = None,
        groups: Optional[Dict[str, Any]] = None,
    ) -> None:
        if not CREWAI_AVAILABLE:
            from ..exceptions import ProviderError

            raise ProviderError(
                "CrewAI SDK not installed. Install with: pip install crewai"
            )

        ctx = get_active_context()

        self.amplitude = amplitude
        self.user_id = user_id or (ctx.user_id if ctx else None) or ""
        self.privacy_config = privacy_config or PrivacyConfig()
        self.session_id = session_id or (ctx.session_id if ctx else None) or str(uuid.uuid4())
        self.trace_id = trace_id or (ctx.trace_id if ctx else None) or str(uuid.uuid4())
        self.agent_id = agent_id or (ctx.agent_id if ctx else None)
        self.parent_agent_id = parent_agent_id or (ctx.parent_agent_id if ctx else None)
        self.customer_org_id = customer_org_id or (ctx.customer_org_id if ctx else None)
        self.agent_version = agent_version or (ctx.agent_version if ctx else None)
        self.context = context or (ctx.context if ctx else None)
        self.env = env or (ctx.env if ctx else None)
        self.groups = groups or (ctx.groups if ctx else None)

        self._turn_counter = 1
        self._llm_start_times: Dict[str, float] = {}
        self._tool_start_times: Dict[str, float] = {}
        self._active = False
        self._listener: Any = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def activate(self) -> None:
        """Register hooks on the CrewAI event bus."""
        if self._active:
            return

        if not CREWAI_AVAILABLE:
            return

        self._listener = _create_listener(self)
        self._active = True

    def deactivate(self) -> None:
        """Unregister hooks from the CrewAI event bus."""
        self._active = False
        self._listener = None

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    def __enter__(self) -> AmplitudeCrewAIHooks:
        self.activate()
        return self

    def __exit__(self, *exc: Any) -> None:
        self.deactivate()

    # ------------------------------------------------------------------
    # Event handlers
    # ------------------------------------------------------------------

    def _on_llm_call_started(self, source: Any, event: Any) -> None:
        """Handle LLMCallStartedEvent."""
        if not self._active:
            return
        call_id = _get_event_id(event)
        self._llm_start_times[call_id] = time.time()

    def _on_llm_call_completed(self, source: Any, event: Any) -> None:
        """Handle LLMCallCompletedEvent -> [Agent] AI Response."""
        if not self._active:
            return

        _logger = get_logger(self.amplitude)

        try:
            call_id = _get_event_id(event)
            start_time = self._llm_start_times.pop(call_id, time.time())
            latency_ms = (time.time() - start_time) * 1000

            response = getattr(event, "response", None)
            response_text = str(response) if response else ""

            model = getattr(event, "model", None) or "unknown"
            provider = infer_provider_from_model(model)

            token_usage = getattr(event, "token_usage", None) or {}
            input_tokens = _safe_int(token_usage.get("prompt_tokens") if isinstance(token_usage, dict)
                                     else getattr(token_usage, "prompt_tokens", None))
            output_tokens = _safe_int(token_usage.get("completion_tokens") if isinstance(token_usage, dict)
                                      else getattr(token_usage, "completion_tokens", None))
            total_tokens = _safe_int(token_usage.get("total_tokens") if isinstance(token_usage, dict)
                                     else getattr(token_usage, "total_tokens", None))

            total_cost_usd = None
            if input_tokens is not None and output_tokens is not None:
                try:
                    total_cost_usd = calculate_cost(
                        model_name=model,
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        default_provider=infer_provider(model),
                    )
                except Exception:
                    pass

            crewai_agent = getattr(event, "agent", None)
            effective_agent_id = self.agent_id
            if crewai_agent:
                agent_role = getattr(crewai_agent, "role", None)
                if agent_role and not effective_agent_id:
                    effective_agent_id = str(agent_role)

            track_ai_message(
                amplitude=self.amplitude,
                user_id=self.user_id,
                model_name=model,
                provider=provider,
                response_content=response_text,
                latency_ms=latency_ms,
                session_id=self.session_id,
                trace_id=self.trace_id,
                turn_id=self._turn_counter,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                total_cost_usd=total_cost_usd,
                agent_id=effective_agent_id,
                parent_agent_id=self.parent_agent_id,
                customer_org_id=self.customer_org_id,
                agent_version=self.agent_version,
                context=self.context,
                env=self.env,
                groups=self.groups,
                privacy_config=self.privacy_config,
            )
            self._turn_counter += 1
        except Exception as exc:
            _logger.warning("CrewAI LLM tracking failed: %s", exc)

    def _on_tool_usage_started(self, source: Any, event: Any) -> None:
        """Handle ToolUsageStartedEvent."""
        if not self._active:
            return
        call_id = _get_event_id(event)
        self._tool_start_times[call_id] = time.time()

    def _on_tool_usage_finished(self, source: Any, event: Any) -> None:
        """Handle ToolUsageFinishedEvent -> [Agent] Tool Call."""
        if not self._active:
            return

        _logger = get_logger(self.amplitude)

        try:
            call_id = _get_event_id(event)
            start_time = self._tool_start_times.pop(call_id, time.time())
            latency_ms = (time.time() - start_time) * 1000

            tool_name = getattr(event, "tool_name", None) or "unknown"
            tool_output = getattr(event, "output", None)
            tool_input_raw = getattr(event, "tool_args", None) or getattr(event, "arguments", None)

            tool_input = None
            if isinstance(tool_input_raw, dict):
                tool_input = tool_input_raw
            elif tool_input_raw is not None:
                tool_input = {"raw": str(tool_input_raw)}

            error = getattr(event, "error", None)
            success = error is None

            crewai_agent = getattr(event, "agent", None)
            effective_agent_id = self.agent_id
            if crewai_agent:
                agent_role = getattr(crewai_agent, "role", None)
                if agent_role and not effective_agent_id:
                    effective_agent_id = str(agent_role)

            track_tool_call(
                amplitude=self.amplitude,
                user_id=self.user_id,
                tool_name=tool_name,
                success=success,
                latency_ms=latency_ms,
                session_id=self.session_id,
                trace_id=self.trace_id,
                turn_id=self._turn_counter,
                tool_input=tool_input,
                tool_output=str(tool_output) if tool_output is not None else None,
                error_message=str(error) if error else None,
                agent_id=effective_agent_id,
                parent_agent_id=self.parent_agent_id,
                customer_org_id=self.customer_org_id,
                agent_version=self.agent_version,
                context=self.context,
                env=self.env,
                groups=self.groups,
                privacy_config=self.privacy_config,
            )
        except Exception as exc:
            _logger.warning("CrewAI tool tracking failed: %s", exc)


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------


def _create_listener(hooks: AmplitudeCrewAIHooks) -> Any:
    """Create and register a CrewAI EventListener that delegates to hooks.

    CrewAI uses a decorator-based event system; we create an EventListener
    class dynamically and register handlers on it.
    """
    try:
        from crewai.utilities.events.event_listener import EventListener
        from crewai.utilities.events.llm_events import (
            LLMCallStartedEvent,
            LLMCallCompletedEvent,
        )
        from crewai.utilities.events.tool_events import (
            ToolUsageStartedEvent,
            ToolUsageFinishedEvent,
        )

        listener = EventListener()

        @listener.on(LLMCallStartedEvent)  # type: ignore[misc]
        def _llm_started(source: Any, event: Any) -> None:
            hooks._on_llm_call_started(source, event)

        @listener.on(LLMCallCompletedEvent)  # type: ignore[misc]
        def _llm_completed(source: Any, event: Any) -> None:
            hooks._on_llm_call_completed(source, event)

        @listener.on(ToolUsageStartedEvent)  # type: ignore[misc]
        def _tool_started(source: Any, event: Any) -> None:
            hooks._on_tool_usage_started(source, event)

        @listener.on(ToolUsageFinishedEvent)  # type: ignore[misc]
        def _tool_finished(source: Any, event: Any) -> None:
            hooks._on_tool_usage_finished(source, event)

        return listener
    except Exception:
        return None


def _get_event_id(event: Any) -> str:
    """Extract a stable ID from a CrewAI event for correlating start/end."""
    eid = getattr(event, "id", None) or getattr(event, "call_id", None)
    return str(eid) if eid else str(id(event))


def _safe_int(val: Any) -> Optional[int]:
    """Convert a value to int, returning None if not possible."""
    if val is None:
        return None
    try:
        return int(val)
    except (ValueError, TypeError):
        return None


