"""CrewAI auto-instrumentation.

Patches ``crewai.Crew.kickoff`` and ``kickoff_async`` to capture
agent_planning spans for the overall crew execution, plus individual
task spans.
"""
from __future__ import annotations

import functools
import logging
from datetime import datetime, timezone
from typing import Any

from fibinaci.instrumentation._base import BaseInstrumentor
from fibinaci.instrumentation._wrapper import ProviderConfig, _make_error_span

logger = logging.getLogger("fibinaci.instrumentation")

_CONFIG = ProviderConfig(
    name="crewai",
    input_tokens_field="prompt_tokens",
    output_tokens_field="completion_tokens",
    finish_reason_getter=lambda r: None,
)


def _extract_crew_info(self_ref: Any) -> dict:
    """Extract metadata from a CrewAI Crew instance."""
    agents = getattr(self_ref, "agents", []) or []
    tasks = getattr(self_ref, "tasks", []) or []
    return {
        "agent_count": len(agents),
        "task_count": len(tasks),
        "process": str(getattr(self_ref, "process", "sequential")),
        "agents": [getattr(a, "role", str(a)) for a in agents[:10]],
    }


def _wrap_kickoff(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData

                crew_info = _extract_crew_info(self)
                # Extract token usage from crew result if available
                usage = getattr(response, "token_usage", None) or {}
                inp = usage.get("prompt_tokens") if isinstance(usage, dict) else getattr(usage, "prompt_tokens", None)
                out = usage.get("completion_tokens") if isinstance(usage, dict) else getattr(usage, "completion_tokens", None)

                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="agent_planning",
                    provider="crewai",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                )
                sd.agent_planning = crew_info
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture crewai span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_make_error_span(_CONFIG, {}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture crewai error span", exc_info=True)
            raise

    return wrapper


def _wrap_akickoff(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData

                crew_info = _extract_crew_info(self)
                usage = getattr(response, "token_usage", None) or {}
                inp = usage.get("prompt_tokens") if isinstance(usage, dict) else getattr(usage, "prompt_tokens", None)
                out = usage.get("completion_tokens") if isinstance(usage, dict) else getattr(usage, "completion_tokens", None)

                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="agent_planning",
                    provider="crewai",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                )
                sd.agent_planning = crew_info
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async crewai span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_make_error_span(_CONFIG, {}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async crewai error span", exc_info=True)
            raise

    return wrapper


class CrewAIInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._original_kickoff: Any = None
        self._original_akickoff: Any = None

    def instrument(self) -> None:
        from crewai import Crew

        self._original_kickoff = Crew.kickoff
        Crew.kickoff = _wrap_kickoff(self._original_kickoff)

        if hasattr(Crew, "kickoff_async"):
            self._original_akickoff = Crew.kickoff_async
            Crew.kickoff_async = _wrap_akickoff(self._original_akickoff)

    def uninstrument(self) -> None:
        if self._original_kickoff is None:
            return
        from crewai import Crew

        Crew.kickoff = self._original_kickoff
        if self._original_akickoff is not None:
            Crew.kickoff_async = self._original_akickoff
        self._original_kickoff = None
        self._original_akickoff = None
