"""CrewAI integration for AgentTrace.

Usage::

    from agenttrace.crewai import AgentTraceCrewAIHandler

    handler = AgentTraceCrewAIHandler()
    crew = Crew(
        ...,
        step_callback=handler.step_callback,
        task_callback=handler.task_callback,
    )
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

from agenttrace.trace import trace


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _safe_str(value: Any) -> str:
    try:
        return str(value)
    except Exception:  # noqa: BLE001
        return "<unserializable>"


class AgentTraceCrewAIHandler:
    """Callback handler for CrewAI step_callback / task_callback hooks."""

    def step_callback(self, step_output: Any) -> None:
        """Called by CrewAI after each agent step."""
        ctx = trace.current()
        if ctx is None:
            return

        now = _now_utc()
        text = _safe_str(step_output)

        ctx.add_span(
            span_type="agent",
            name="crewai_step",
            start_time=now,
            end_time=now,
            input_json={},
            output_json={"step_output": text},
            metadata_json={"source": "crewai"},
        )

    def task_callback(self, task_output: Any) -> None:
        """Called by CrewAI after each task completes."""
        ctx = trace.current()
        if ctx is None:
            return

        now = _now_utc()
        output_data: dict[str, Any] = {"task_output": _safe_str(task_output)}

        # CrewAI TaskOutput objects may have structured attributes
        for attr in ("description", "summary", "raw"):
            val = getattr(task_output, attr, None)
            if val is not None:
                output_data[attr] = _safe_str(val)

        ctx.add_span(
            span_type="chain",
            name="crewai_task",
            start_time=now,
            end_time=now,
            input_json={},
            output_json=output_data,
            metadata_json={"source": "crewai"},
        )
