"""CrewAI adapter — instrument your crew in 2 lines.

Usage:
    from velenai_sdk import VelenAI
    from velenai_sdk.adapters.crewai import VelenAICrewCallback

    vai = VelenAI()
    crew = Crew(agents=[...], tasks=[...], callbacks=[VelenAICrewCallback(vai)])
"""

from __future__ import annotations

import time
from typing import Any

from ..client import VelenAI


class VelenAICrewCallback:
    """CrewAI callback handler that emits telemetry to VelenAI.

    Tracks each task execution as a separate agent run, mapping:
    - CrewAI task.agent.role → VelenAI agent_id
    - Task success/failure → status
    - Duration → latency_ms
    - Token usage from LLM calls → tokens
    """

    def __init__(self, client: VelenAI, agent_prefix: str = "crewai"):
        self._client = client
        self._prefix = agent_prefix
        self._task_starts: dict[str, float] = {}

    def on_task_start(self, task: Any, **kwargs: Any) -> None:
        task_id = str(getattr(task, "id", id(task)))
        self._task_starts[task_id] = time.monotonic()

    def on_task_end(self, task: Any, output: Any, **kwargs: Any) -> None:
        task_id = str(getattr(task, "id", id(task)))
        start = self._task_starts.pop(task_id, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        agent_role = "unknown"
        if hasattr(task, "agent") and hasattr(task.agent, "role"):
            agent_role = task.agent.role.lower().replace(" ", "_")

        agent_id = f"{self._prefix}.{agent_role}"

        # Extract token usage if available
        tokens = None
        cost_usd = None
        if hasattr(output, "token_usage"):
            usage = output.token_usage
            tokens = getattr(usage, "total_tokens", None)
        if hasattr(output, "cost"):
            cost_usd = getattr(output, "cost", None)

        metadata = {
            "framework": "crewai",
            "task_description": str(getattr(task, "description", ""))[:200],
        }

        self._client.emit(
            agent_id=agent_id,
            success=True,
            latency_ms=latency_ms,
            tokens=tokens,
            cost_usd=cost_usd,
            metadata=metadata,
        )

    def on_task_error(self, task: Any, error: Exception, **kwargs: Any) -> None:
        task_id = str(getattr(task, "id", id(task)))
        start = self._task_starts.pop(task_id, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        agent_role = "unknown"
        if hasattr(task, "agent") and hasattr(task.agent, "role"):
            agent_role = task.agent.role.lower().replace(" ", "_")

        agent_id = f"{self._prefix}.{agent_role}"

        self._client.emit(
            agent_id=agent_id,
            success=False,
            latency_ms=latency_ms,
            metadata={
                "framework": "crewai",
                "error": str(error)[:500],
            },
        )

    def on_crew_start(self, crew: Any, **kwargs: Any) -> None:
        """Optional: track the entire crew run."""
        self._crew_start = time.monotonic()

    def on_crew_end(self, crew: Any, output: Any, **kwargs: Any) -> None:
        """Optional: emit a crew-level event."""
        if hasattr(self, "_crew_start"):
            latency_ms = int((time.monotonic() - self._crew_start) * 1000)
            crew_name = getattr(crew, "name", "crew")
            self._client.emit(
                agent_id=f"{self._prefix}.{crew_name}",
                success=True,
                latency_ms=latency_ms,
                metadata={"framework": "crewai", "type": "crew_run"},
            )
