"""
CrewAI integration for Visibe.

Uses OpenTelemetry to capture metrics from CrewAI 1.8.0+
"""

import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
import uuid

from opentelemetry import trace

# SDK logger
logger = logging.getLogger("visibe")
from opentelemetry.sdk.trace import TracerProvider, ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult, SimpleSpanProcessor
from opentelemetry.trace import StatusCode

from .base import BaseIntegration, TraceSummary
from ..utils import calculate_cost

# Check if CrewAI instrumentation is available
try:
    from openinference.instrumentation.crewai import CrewAIInstrumentor
    from openinference.instrumentation.openai import OpenAIInstrumentor
    CREWAI_INSTRUMENTATION_AVAILABLE = True
except ImportError:
    CREWAI_INSTRUMENTATION_AVAILABLE = False


# Internal tool span names to filter out
INTERNAL_TOOL_NAMES = {"Tool Usage Error", "_internal"}

# Tracks instrumented crews: {id(crew): {'kickoff': original_method, 'name': str}}
_instrumented_crews: Dict[int, Dict[str, Any]] = {}

# Module-level flag: set to True while a CrewAI trace is active.
# The OpenAI patcher checks this to avoid creating duplicate standalone traces.
_active_crew_trace = False


class CrewAIIntegration(BaseIntegration):
    """
    CrewAI integration for Visibe observability.

    Usage:
        from visibe import Visibe
        from visibe.integrations.crewai import CrewAIIntegration

        obs = Visibe(api_url="http://localhost:3000")
        tracker = CrewAIIntegration(obs)

        crew = Crew(agents=[...], tasks=[...])
        with tracker.track(crew, name="my-workflow"):
            result = crew.kickoff()
    """

    _instrumentation_setup = False

    def __init__(self, client):
        """
        Initialize CrewAI integration.

        Args:
            client: Visibe client instance
        """
        if not CREWAI_INSTRUMENTATION_AVAILABLE:
            raise ImportError(
                "CrewAI integration requires OpenInference instrumentation.\n"
                "Install with: pip install \"visibe[crewai]\"\n"
            )
        super().__init__(client)

        # Set up OpenTelemetry instrumentation (only once globally)
        if not CrewAIIntegration._instrumentation_setup:
            self._setup_instrumentation()
            CrewAIIntegration._instrumentation_setup = True

    def _setup_instrumentation(self) -> None:
        """
        Set up OpenTelemetry instrumentation for CrewAI.
        """
        # Get the current tracer provider
        existing_provider = trace.get_tracer_provider()

        # Check if it's a real SDK provider or just a proxy
        provider_type = type(existing_provider).__name__
        logger.debug(f"Current TracerProvider type: {provider_type}")

        # If it's a ProxyTracerProvider, we need to create a real one
        if provider_type == 'ProxyTracerProvider' or not hasattr(existing_provider, 'add_span_processor'):
            logger.debug("Creating new SDK TracerProvider")
            tracer_provider = TracerProvider()
            trace.set_tracer_provider(tracer_provider)
        else:
            logger.debug("Using existing TracerProvider")
            tracer_provider = existing_provider

        # Instrument CrewAI
        try:
            CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
            logger.debug("CrewAI instrumented successfully")
        except Exception as e:
            logger.warning(f"Failed to instrument CrewAI: {e}")

        # Instrument OpenAI (the LLM provider)
        try:
            OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)
            logger.debug("OpenAI instrumented successfully")
        except Exception as e:
            logger.warning(f"Failed to instrument OpenAI: {e}")

    def _add_span_processor(self, exporter: 'CrewAISpanExporter') -> SimpleSpanProcessor:
        """Add a span processor for the given exporter to the tracer provider.
        
        Returns the SimpleSpanProcessor so callers can shut it down after use,
        preventing processor accumulation across repeated invocations.
        """
        tracer_provider = trace.get_tracer_provider()
        provider_type = type(tracer_provider).__name__

        if not hasattr(tracer_provider, 'add_span_processor'):
            logger.error(
                f"TracerProvider ({provider_type}) doesn't support "
                f"add_span_processor — spans won't be captured"
            )
            return None

        try:
            span_processor = SimpleSpanProcessor(exporter)
            tracer_provider.add_span_processor(span_processor)
            logger.debug("Span processor added successfully")
            return span_processor
        except Exception as e:
            logger.error(f"Failed to add span processor: {e}")
            return None

    @staticmethod
    def _shutdown_span_processor(span_processor: SimpleSpanProcessor) -> None:
        """Shut down a span processor to prevent accumulation."""
        if span_processor is None:
            return
        try:
            span_processor.shutdown()
            logger.debug("Span processor shut down successfully")
        except Exception as e:
            logger.debug(f"Span processor shutdown: {e}")

    def track(self, crew, name: str = "crew-execution"):
        """
        Track a CrewAI crew execution via context manager.

        Args:
            crew: CrewAI Crew instance
            name: Human-readable name for this workflow

        Returns:
            CrewTracker context manager
        """
        tracker = CrewTracker(self.client, crew, name)
        tracker._span_processor = self._add_span_processor(tracker.exporter)
        return tracker

    # ==================== instrument() mode ====================

    def instrument(self, crew, name: Optional[str] = None):
        """Patch crew.kickoff() so every execution sends a trace automatically.

        Args:
            crew: CrewAI Crew instance
            name: Optional name for traces. Shown in the Visibe dashboard.

        Example:
            >>> obs = Visibe()
            >>> obs.instrument(crew, name="my-workflow")
            >>> result = crew.kickoff()  # trace sent automatically
        """
        crew_id = id(crew)
        if crew_id in _instrumented_crews:
            if not name:
                logger.debug("Crew already instrumented — skipping")
                return
            logger.debug(f"Crew already instrumented, updating trace name to: {name}")
            return

        original_kickoff = crew.kickoff
        _instrumented_crews[crew_id] = {
            'kickoff': original_kickoff,
            'name': name,
        }

        integration = self
        trace_name = name

        def _derive_trace_name(crew_obj):
            """Derive a trace name from the crew when none was provided."""
            if trace_name:
                return trace_name
            # Use crew.name if the user set a custom one (default is 'crew')
            crew_name = getattr(crew_obj, 'name', None)
            if crew_name and crew_name != 'crew':
                return crew_name
            # Fall back to a descriptive name from agent roles
            try:
                agents = getattr(crew_obj, 'agents', None)
                if agents:
                    roles = [getattr(a, 'role', None) for a in agents[:3]]
                    roles = [r for r in roles if r]
                    if roles:
                        label = ', '.join(roles)
                        if len(agents) > 3:
                            label += f' +{len(agents) - 3} more'
                        return label
            except Exception:
                pass
            return 'crewai-execution'

        def wrapped_kickoff(*args, **kwargs):
            global _active_crew_trace
            exporter = CrewAISpanExporter()
            exporter.populate_task_descriptions(crew)
            _span_processor = integration._add_span_processor(exporter)

            trace_id = str(uuid.uuid4())
            effective_name = _derive_trace_name(crew)

            # Create trace header
            integration.client.create_trace({
                'trace_id': trace_id,
                'name': effective_name,
                'framework': 'crewai',
                'started_at': datetime.now(timezone.utc).isoformat(),
            })

            # Wire real-time span streaming
            exporter._span_sender = lambda span: integration.client.queue_span(
                trace_id, span
            )

            # Signal to the OpenAI patcher that a CrewAI trace is active
            _active_crew_trace = True

            status = 'completed'
            try:
                result = original_kickoff(*args, **kwargs)
                return result
            except Exception as exc:
                status = 'failed'
                raise
            finally:
                _active_crew_trace = False
                # Force flush OTEL spans
                tracer_provider = trace.get_tracer_provider()
                if hasattr(tracer_provider, 'force_flush'):
                    try:
                        tracer_provider.force_flush(timeout_millis=5000)
                    except Exception as e:
                        logger.warning(f"Force flush failed: {e}")

                metrics = exporter.get_metrics()

                # Extract prompt from first task
                prompt = None
                if hasattr(crew, 'tasks') and crew.tasks:
                    first_task = crew.tasks[0]
                    if hasattr(first_task, 'description'):
                        prompt = first_task.description

                if exporter.errors:
                    status = 'failed'

                total_input_tokens = metrics['total_input_tokens']
                total_output_tokens = metrics['total_output_tokens']
                total_cost = metrics['total_cost']
                total_tokens = total_input_tokens + total_output_tokens
                model = exporter.llm_calls[0]['model'] if exporter.llm_calls else None

                success = integration.client.complete_trace(trace_id, {
                    'status': status,
                    'ended_at': metrics['end_time'],
                    'duration_ms': metrics['duration_ms'],
                    'prompt': prompt or "No prompt captured",
                    'model': model,
                    'total_cost': total_cost,
                    'total_tokens': total_tokens,
                    'total_input_tokens': total_input_tokens,
                    'total_output_tokens': total_output_tokens,
                })

                duration_s = metrics['duration_ms'] / 1000
                summary = TraceSummary(
                    name=effective_name or 'crewai',
                    status=status,
                    llm_calls=len(exporter.llm_calls),
                    tool_calls=len(exporter.tool_calls),
                    total_tokens=total_tokens,
                    total_cost=total_cost,
                    duration_s=duration_s,
                    agents=exporter.agents_executed,
                    sent=success,
                )
                print(f"[Visibe] {summary}")

                # Clean up span processor to prevent accumulation
                integration._shutdown_span_processor(_span_processor)

        # Crew is a Pydantic model — bypass __setattr__ validation
        object.__setattr__(crew, 'kickoff', wrapped_kickoff)

        # Patch kickoff_async if available
        if hasattr(crew, 'kickoff_async'):
            original_kickoff_async = crew.kickoff_async
            _instrumented_crews[crew_id]['kickoff_async'] = original_kickoff_async

            async def wrapped_kickoff_async(*args, **kwargs):
                global _active_crew_trace
                exporter = CrewAISpanExporter()
                exporter.populate_task_descriptions(crew)
                _span_processor = integration._add_span_processor(exporter)

                trace_id = str(uuid.uuid4())
                effective_name = _derive_trace_name(crew)

                # Create trace header
                integration.client.create_trace({
                    'trace_id': trace_id,
                    'name': effective_name,
                    'framework': 'crewai',
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                # Wire real-time span streaming
                exporter._span_sender = lambda span: integration.client.queue_span(
                    trace_id, span
                )

                # Signal to the OpenAI patcher that a CrewAI trace is active
                _active_crew_trace = True

                status = 'completed'
                try:
                    result = await original_kickoff_async(*args, **kwargs)
                    return result
                except Exception:
                    status = 'failed'
                    raise
                finally:
                    _active_crew_trace = False
                    # Force flush OTEL spans
                    tracer_provider = trace.get_tracer_provider()
                    if hasattr(tracer_provider, 'force_flush'):
                        try:
                            tracer_provider.force_flush(timeout_millis=5000)
                        except Exception as e:
                            logger.warning(f"Force flush failed: {e}")

                    metrics = exporter.get_metrics()

                    # Extract prompt from first task
                    prompt = None
                    if hasattr(crew, 'tasks') and crew.tasks:
                        first_task = crew.tasks[0]
                        if hasattr(first_task, 'description'):
                            prompt = first_task.description

                    if exporter.errors:
                        status = 'failed'

                    total_input_tokens = metrics['total_input_tokens']
                    total_output_tokens = metrics['total_output_tokens']
                    total_cost = metrics['total_cost']
                    total_tokens = total_input_tokens + total_output_tokens
                    model = exporter.llm_calls[0]['model'] if exporter.llm_calls else None

                    success = integration.client.complete_trace(trace_id, {
                        'status': status,
                        'ended_at': metrics['end_time'],
                        'duration_ms': metrics['duration_ms'],
                        'prompt': prompt or "No prompt captured",
                        'model': model,
                        'total_cost': total_cost,
                        'total_tokens': total_tokens,
                        'total_input_tokens': total_input_tokens,
                        'total_output_tokens': total_output_tokens,
                    })

                    duration_s = metrics['duration_ms'] / 1000
                    summary = TraceSummary(
                        name=effective_name or 'crewai',
                        status=status,
                        llm_calls=len(exporter.llm_calls),
                        tool_calls=len(exporter.tool_calls),
                        total_tokens=total_tokens,
                        total_cost=total_cost,
                        duration_s=duration_s,
                        agents=exporter.agents_executed,
                        sent=success,
                    )
                    print(f"[Visibe] {summary}")

                    # Clean up span processor to prevent accumulation
                    integration._shutdown_span_processor(_span_processor)

            object.__setattr__(crew, 'kickoff_async', wrapped_kickoff_async)

        # Patch kickoff_for_each if available
        if hasattr(crew, 'kickoff_for_each'):
            original_kickoff_for_each = crew.kickoff_for_each
            _instrumented_crews[crew_id]['kickoff_for_each'] = original_kickoff_for_each

            def wrapped_kickoff_for_each(*args, **kwargs):
                global _active_crew_trace
                exporter = CrewAISpanExporter()
                exporter.populate_task_descriptions(crew)
                _span_processor = integration._add_span_processor(exporter)

                trace_id = str(uuid.uuid4())
                effective_name = _derive_trace_name(crew)

                integration.client.create_trace({
                    'trace_id': trace_id,
                    'name': effective_name,
                    'framework': 'crewai',
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                exporter._span_sender = lambda span: integration.client.queue_span(
                    trace_id, span
                )

                _active_crew_trace = True
                status = 'completed'
                try:
                    result = original_kickoff_for_each(*args, **kwargs)
                    return result
                except Exception:
                    status = 'failed'
                    raise
                finally:
                    _active_crew_trace = False
                    tracer_provider = trace.get_tracer_provider()
                    if hasattr(tracer_provider, 'force_flush'):
                        try:
                            tracer_provider.force_flush(timeout_millis=5000)
                        except Exception as e:
                            logger.warning(f"Force flush failed: {e}")

                    metrics = exporter.get_metrics()

                    prompt = None
                    if hasattr(crew, 'tasks') and crew.tasks:
                        first_task = crew.tasks[0]
                        if hasattr(first_task, 'description'):
                            prompt = first_task.description

                    if exporter.errors:
                        status = 'failed'

                    total_input_tokens = metrics['total_input_tokens']
                    total_output_tokens = metrics['total_output_tokens']
                    total_cost = metrics['total_cost']
                    total_tokens = total_input_tokens + total_output_tokens
                    model = exporter.llm_calls[0]['model'] if exporter.llm_calls else None

                    success = integration.client.complete_trace(trace_id, {
                        'status': status,
                        'ended_at': metrics['end_time'],
                        'duration_ms': metrics['duration_ms'],
                        'prompt': prompt or "No prompt captured",
                        'model': model,
                        'total_cost': total_cost,
                        'total_tokens': total_tokens,
                        'total_input_tokens': total_input_tokens,
                        'total_output_tokens': total_output_tokens,
                    })

                    duration_s = metrics['duration_ms'] / 1000
                    summary = TraceSummary(
                        name=effective_name or 'crewai',
                        status=status,
                        llm_calls=len(exporter.llm_calls),
                        tool_calls=len(exporter.tool_calls),
                        total_tokens=total_tokens,
                        total_cost=total_cost,
                        duration_s=duration_s,
                        agents=exporter.agents_executed,
                        sent=success,
                    )
                    print(f"[Visibe] {summary}")

                    # Clean up span processor to prevent accumulation
                    integration._shutdown_span_processor(_span_processor)

            object.__setattr__(crew, 'kickoff_for_each', wrapped_kickoff_for_each)

        # Patch kickoff_for_each_async if available
        if hasattr(crew, 'kickoff_for_each_async'):
            original_kickoff_for_each_async = crew.kickoff_for_each_async
            _instrumented_crews[crew_id]['kickoff_for_each_async'] = original_kickoff_for_each_async

            async def wrapped_kickoff_for_each_async(*args, **kwargs):
                global _active_crew_trace
                exporter = CrewAISpanExporter()
                exporter.populate_task_descriptions(crew)
                _span_processor = integration._add_span_processor(exporter)

                trace_id = str(uuid.uuid4())
                effective_name = _derive_trace_name(crew)

                integration.client.create_trace({
                    'trace_id': trace_id,
                    'name': effective_name,
                    'framework': 'crewai',
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                exporter._span_sender = lambda span: integration.client.queue_span(
                    trace_id, span
                )

                _active_crew_trace = True
                status = 'completed'
                try:
                    result = await original_kickoff_for_each_async(*args, **kwargs)
                    return result
                except Exception:
                    status = 'failed'
                    raise
                finally:
                    _active_crew_trace = False
                    tracer_provider = trace.get_tracer_provider()
                    if hasattr(tracer_provider, 'force_flush'):
                        try:
                            tracer_provider.force_flush(timeout_millis=5000)
                        except Exception as e:
                            logger.warning(f"Force flush failed: {e}")

                    metrics = exporter.get_metrics()

                    prompt = None
                    if hasattr(crew, 'tasks') and crew.tasks:
                        first_task = crew.tasks[0]
                        if hasattr(first_task, 'description'):
                            prompt = first_task.description

                    if exporter.errors:
                        status = 'failed'

                    total_input_tokens = metrics['total_input_tokens']
                    total_output_tokens = metrics['total_output_tokens']
                    total_cost = metrics['total_cost']
                    total_tokens = total_input_tokens + total_output_tokens
                    model = exporter.llm_calls[0]['model'] if exporter.llm_calls else None

                    success = integration.client.complete_trace(trace_id, {
                        'status': status,
                        'ended_at': metrics['end_time'],
                        'duration_ms': metrics['duration_ms'],
                        'prompt': prompt or "No prompt captured",
                        'model': model,
                        'total_cost': total_cost,
                        'total_tokens': total_tokens,
                        'total_input_tokens': total_input_tokens,
                        'total_output_tokens': total_output_tokens,
                    })

                    duration_s = metrics['duration_ms'] / 1000
                    summary = TraceSummary(
                        name=effective_name or 'crewai',
                        status=status,
                        llm_calls=len(exporter.llm_calls),
                        tool_calls=len(exporter.tool_calls),
                        total_tokens=total_tokens,
                        total_cost=total_cost,
                        duration_s=duration_s,
                        agents=exporter.agents_executed,
                        sent=success,
                    )
                    print(f"[Visibe] {summary}")

                    # Clean up span processor to prevent accumulation
                    integration._shutdown_span_processor(_span_processor)

            object.__setattr__(crew, 'kickoff_for_each_async', wrapped_kickoff_for_each_async)

        # Patch train if available
        if hasattr(crew, 'train'):
            original_train = crew.train
            _instrumented_crews[crew_id]['train'] = original_train

            def wrapped_train(*args, **kwargs):
                exporter = CrewAISpanExporter()
                exporter.populate_task_descriptions(crew)
                _span_processor = integration._add_span_processor(exporter)

                trace_id = str(uuid.uuid4())

                integration.client.create_trace({
                    'trace_id': trace_id,
                    'name': trace_name or "crew-training",
                    'framework': 'crewai',
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                exporter._span_sender = lambda span: integration.client.queue_span(
                    trace_id, span
                )

                status = 'completed'
                try:
                    result = original_train(*args, **kwargs)
                    return result
                except Exception:
                    status = 'failed'
                    raise
                finally:
                    tracer_provider = trace.get_tracer_provider()
                    if hasattr(tracer_provider, 'force_flush'):
                        try:
                            tracer_provider.force_flush(timeout_millis=5000)
                        except Exception as e:
                            logger.warning(f"Force flush failed: {e}")

                    metrics = exporter.get_metrics()

                    prompt = None
                    if hasattr(crew, 'tasks') and crew.tasks:
                        first_task = crew.tasks[0]
                        if hasattr(first_task, 'description'):
                            prompt = first_task.description

                    if exporter.errors:
                        status = 'failed'

                    total_input_tokens = metrics['total_input_tokens']
                    total_output_tokens = metrics['total_output_tokens']
                    total_cost = metrics['total_cost']
                    total_tokens = total_input_tokens + total_output_tokens
                    model = exporter.llm_calls[0]['model'] if exporter.llm_calls else None

                    success = integration.client.complete_trace(trace_id, {
                        'status': status,
                        'ended_at': metrics['end_time'],
                        'duration_ms': metrics['duration_ms'],
                        'prompt': prompt or "No prompt captured",
                        'model': model,
                        'total_cost': total_cost,
                        'total_tokens': total_tokens,
                        'total_input_tokens': total_input_tokens,
                        'total_output_tokens': total_output_tokens,
                    })

                    duration_s = metrics['duration_ms'] / 1000
                    summary = TraceSummary(
                        name=trace_name or 'crew-training',
                        status=status,
                        llm_calls=len(exporter.llm_calls),
                        tool_calls=len(exporter.tool_calls),
                        total_tokens=total_tokens,
                        total_cost=total_cost,
                        duration_s=duration_s,
                        agents=exporter.agents_executed,
                        sent=success,
                    )
                    print(f"[Visibe] {summary}")

                    # Clean up span processor to prevent accumulation
                    integration._shutdown_span_processor(_span_processor)

            object.__setattr__(crew, 'train', wrapped_train)

        # Patch test if available
        if hasattr(crew, 'test'):
            original_test = crew.test
            _instrumented_crews[crew_id]['test'] = original_test

            def wrapped_test(*args, **kwargs):
                exporter = CrewAISpanExporter()
                exporter.populate_task_descriptions(crew)
                _span_processor = integration._add_span_processor(exporter)

                trace_id = str(uuid.uuid4())

                integration.client.create_trace({
                    'trace_id': trace_id,
                    'name': trace_name or "crew-test",
                    'framework': 'crewai',
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                exporter._span_sender = lambda span: integration.client.queue_span(
                    trace_id, span
                )

                status = 'completed'
                try:
                    result = original_test(*args, **kwargs)
                    return result
                except Exception:
                    status = 'failed'
                    raise
                finally:
                    tracer_provider = trace.get_tracer_provider()
                    if hasattr(tracer_provider, 'force_flush'):
                        try:
                            tracer_provider.force_flush(timeout_millis=5000)
                        except Exception as e:
                            logger.warning(f"Force flush failed: {e}")

                    metrics = exporter.get_metrics()

                    prompt = None
                    if hasattr(crew, 'tasks') and crew.tasks:
                        first_task = crew.tasks[0]
                        if hasattr(first_task, 'description'):
                            prompt = first_task.description

                    if exporter.errors:
                        status = 'failed'

                    total_input_tokens = metrics['total_input_tokens']
                    total_output_tokens = metrics['total_output_tokens']
                    total_cost = metrics['total_cost']
                    total_tokens = total_input_tokens + total_output_tokens
                    model = exporter.llm_calls[0]['model'] if exporter.llm_calls else None

                    success = integration.client.complete_trace(trace_id, {
                        'status': status,
                        'ended_at': metrics['end_time'],
                        'duration_ms': metrics['duration_ms'],
                        'prompt': prompt or "No prompt captured",
                        'model': model,
                        'total_cost': total_cost,
                        'total_tokens': total_tokens,
                        'total_input_tokens': total_input_tokens,
                        'total_output_tokens': total_output_tokens,
                    })

                    duration_s = metrics['duration_ms'] / 1000
                    summary = TraceSummary(
                        name=trace_name or 'crew-test',
                        status=status,
                        llm_calls=len(exporter.llm_calls),
                        tool_calls=len(exporter.tool_calls),
                        total_tokens=total_tokens,
                        total_cost=total_cost,
                        duration_s=duration_s,
                        agents=exporter.agents_executed,
                        sent=success,
                    )
                    print(f"[Visibe] {summary}")

                    # Clean up span processor to prevent accumulation
                    integration._shutdown_span_processor(_span_processor)

            object.__setattr__(crew, 'test', wrapped_test)

        logger.debug(f"Instrumented Crew (id={crew_id})")

    def uninstrument(self, crew):
        """Remove instrumentation from a previously instrumented Crew.

        Args:
            crew: A previously instrumented CrewAI Crew instance
        """
        crew_id = id(crew)
        if crew_id not in _instrumented_crews:
            logger.warning("Crew not instrumented — nothing to restore")
            return

        original = _instrumented_crews.pop(crew_id)
        object.__setattr__(crew, 'kickoff', original['kickoff'])
        if original.get('kickoff_async') is not None:
            object.__setattr__(crew, 'kickoff_async', original['kickoff_async'])
        if original.get('kickoff_for_each') is not None:
            object.__setattr__(crew, 'kickoff_for_each', original['kickoff_for_each'])
        if original.get('kickoff_for_each_async') is not None:
            object.__setattr__(crew, 'kickoff_for_each_async', original['kickoff_for_each_async'])
        if original.get('train') is not None:
            object.__setattr__(crew, 'train', original['train'])
        if original.get('test') is not None:
            object.__setattr__(crew, 'test', original['test'])
        logger.debug(f"Uninstrumented Crew (id={crew_id})")


class CrewAISpanExporter(SpanExporter):
    """
    Custom OpenTelemetry Span Exporter that captures metrics from CrewAI execution.

    This exporter intercepts all OTEL spans generated by CrewAI and extracts:
    - LLM calls (model, tokens, cost)
    - Agent execution (names, timing)
    - Tool usage
    - Errors
    """

    def __init__(self):
        """Initialize the span exporter"""
        self.llm_calls: List[Dict[str, Any]] = []
        self.tool_calls: List[Dict[str, Any]] = []
        self.agents_executed: List[str] = []
        self.agent_timings: Dict[str, Dict[str, Any]] = {}
        self.errors: List[Dict[str, Any]] = []
        self.start_time: Optional[datetime] = None
        self.end_time: Optional[datetime] = None

        # Chronological spans timeline
        self.steps: List[Dict[str, Any]] = []
        self.step_counter = 0

        # Callable to stream spans to backend — set by CrewTracker
        self._span_sender = None

        # Track current agent context for associating LLM/tool calls with agents
        self.current_agent: Optional[str] = None
        self._current_agent_span_id: Optional[str] = None

        # Track current task context for per-task cost breakdown
        self.current_task: Optional[str] = None
        self.task_info: Dict[str, Dict[str, Any]] = {}

    def populate_task_descriptions(self, crew) -> None:
        """Pre-populate task descriptions from Crew task objects.

        Uses task.name if set, otherwise task.description (full text).
        """
        if not hasattr(crew, 'tasks') or not crew.tasks:
            return
        for task in crew.tasks:
            task_id = str(getattr(task, 'id', ''))
            if not task_id:
                continue
            description = getattr(task, 'name', None) or getattr(task, 'description', None)
            if task_id not in self.task_info:
                self.task_info[task_id] = {
                    'task_id': task_id,
                    'description': description,
                    'agent_name': None,
                    'input_tokens': 0,
                    'output_tokens': 0,
                    'cost': 0.0
                }
            elif description and not self.task_info[task_id].get('description'):
                self.task_info[task_id]['description'] = description

    @staticmethod
    def _detect_provider(model: str) -> str:
        """Detect LLM provider from model name."""
        model_lower = model.lower()
        if 'claude' in model_lower or 'anthropic' in model_lower:
            return 'anthropic'
        if 'gemini' in model_lower or 'palm' in model_lower:
            return 'google'
        return 'openai'

    def _create_step(self, step_type: str, timestamp: datetime, **kwargs) -> str:
        """
        Create a span, add it to the timeline, and stream to backend.

        Args:
            step_type: Type of span (e.g., "agent_start", "llm_call", "tool_call")
            timestamp: When this span occurred
            **kwargs: Additional span-specific data

        Returns:
            span_id: Unique identifier for this span
        """
        self.step_counter += 1
        span_id = f'span_{self.step_counter}'

        span = {
            'span_id': span_id,
            'type': step_type,
            'timestamp': timestamp.isoformat(),
            **kwargs
        }

        self.steps.append(span)

        # Stream span to backend in real-time
        if self._span_sender:
            try:
                self._span_sender(span)
            except Exception as e:
                logger.warning(f"Failed to send span: {e}")

        return span_id

    def export(self, spans: List[ReadableSpan]) -> SpanExportResult:
        """
        Process exported spans and extract metrics.

        Args:
            spans: List of OTEL spans from CrewAI execution

        Returns:
            SpanExportResult.SUCCESS
        """
        logger.debug(f"export() called with {len(spans)} span(s)")

        for span in spans:
            logger.debug(f"Processing span: {span.name}")
            self._process_span(span)

        return SpanExportResult.SUCCESS

    def _process_span(self, span: ReadableSpan) -> None:
        """
        Extract metrics from a single span.

        Args:
            span: OTEL span to process
        """
        # Track overall timing
        if self.start_time is None:
            self.start_time = datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc)
        self.end_time = datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc)

        attributes = dict(span.attributes) if span.attributes else {}
        span_name = span.name

        # Log span attributes for debugging
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Span '{span_name}' has {len(attributes)} attributes")
            for key, value in list(attributes.items())[:10]:
                logger.debug(f"  {key} = {value}")

        # Check for task context first (sets current_task for subsequent spans)
        if self._is_task_span(span_name, attributes):
            self._process_task(span, attributes)

        # Check if this is an LLM call span
        if self._is_llm_span(span_name, attributes):
            self._process_llm_call(span, attributes)

        # Check if this is a tool call span
        elif self._is_tool_span(span_name, attributes):
            self._process_tool_call(span, attributes)

        # Check if this is an agent execution span
        elif self._is_agent_span(span_name, attributes):
            self._process_agent_execution(span, attributes)

        # Check for errors in any span
        if span.status.status_code != StatusCode.OK:
            self._process_error(span, attributes)

    def _is_llm_span(self, span_name: str, attributes: Dict[str, Any]) -> bool:
        """Check if span represents an LLM call"""
        return (
            'llm' in span_name.lower() or
            'llm.model_name' in attributes or
            'gen_ai.request.model' in attributes
        )

    def _is_tool_span(self, span_name: str, attributes: Dict[str, Any]) -> bool:
        """Check if span represents a tool call"""
        return (
            'tool' in span_name.lower() or
            'crewai.tool.name' in attributes or
            'tool.name' in attributes or
            attributes.get('openinference.span.kind') == 'TOOL'
        )

    def _is_agent_span(self, span_name: str, attributes: Dict[str, Any]) -> bool:
        """Check if span represents agent execution"""
        return (
            'agent' in span_name.lower() or
            'agent_role' in attributes or
            'graph.node.id' in attributes or
            attributes.get('openinference.span.kind') == 'AGENT' or
            'crewai.agent.name' in attributes or
            'crewai.agent.role' in attributes
        )

    def _is_task_span(self, span_name: str, attributes: Dict[str, Any]) -> bool:
        """Check if span represents a task"""
        return (
            'task' in span_name.lower() or
            'task_id' in attributes or
            'task_key' in attributes
        )

    def _process_task(self, span: ReadableSpan, attributes: Dict[str, Any]) -> None:
        """Extract task information and set current task context"""
        import json

        task_id = attributes.get('task_id') or attributes.get('task_key')
        if not task_id:
            return

        # Set current task context
        self.current_task = task_id

        # Extract task description from output.value if available
        task_description = None
        if 'output.value' in attributes:
            try:
                output_data = json.loads(attributes['output.value'])
                task_description = output_data.get('description') or output_data.get('name')
            except (json.JSONDecodeError, TypeError):
                pass

        # Get agent for this task
        agent_name = attributes.get('agent_role') or self.current_agent

        # Store/update task info
        if task_id not in self.task_info:
            self.task_info[task_id] = {
                'task_id': task_id,
                'description': task_description,
                'agent_name': agent_name,
                'input_tokens': 0,
                'output_tokens': 0,
                'cost': 0.0
            }
        elif task_description and not self.task_info[task_id].get('description'):
            self.task_info[task_id]['description'] = task_description

        logger.debug(f"Task context set: {task_id[:20]}...")

    def _process_llm_call(self, span: ReadableSpan, attributes: Dict[str, Any]) -> None:
        """Extract LLM call metrics from span"""
        import json

        # Extract model name - it's in llm.invocation_parameters as JSON
        model = 'unknown'
        if 'llm.invocation_parameters' in attributes:
            try:
                params = json.loads(attributes['llm.invocation_parameters'])
                model = params.get('model', 'unknown')
            except:
                pass

        # Fallback to other possible attributes
        if model == 'unknown':
            model = (
                attributes.get('llm.model_name') or
                attributes.get('gen_ai.request.model') or
                attributes.get('gen_ai.system') or
                'unknown'
            )

        # Extract token usage - it's in output.value as JSON
        input_tokens = 0
        output_tokens = 0
        if 'output.value' in attributes:
            try:
                output_data = json.loads(attributes['output.value'])
                usage = output_data.get('usage', {})
                input_tokens = usage.get('prompt_tokens', 0)
                output_tokens = usage.get('completion_tokens', 0)
            except:
                pass

        # Fallback to direct attributes (in case format changes)
        if input_tokens == 0:
            input_tokens = (
                attributes.get('llm.usage.prompt_tokens') or
                attributes.get('gen_ai.usage.prompt_tokens') or
                0
            )
        if output_tokens == 0:
            output_tokens = (
                attributes.get('llm.usage.completion_tokens') or
                attributes.get('gen_ai.usage.completion_tokens') or
                0
            )

        total_tokens = input_tokens + output_tokens

        logger.debug(f"Extracted: model={model}, input={input_tokens}, output={output_tokens}")

        # Extract agent name from current context (set by _process_agent_execution)
        agent_name = self.current_agent or 'unknown'

        logger.debug(f"Agent: {agent_name} (from current context)")

        # Extract input/output text from message arrays
        llm_input = ''
        llm_output = ''

        # Input messages - concatenate all user/system messages
        input_messages = []
        i = 0
        while f'llm.input_messages.{i}.message.content' in attributes:
            role = attributes.get(f'llm.input_messages.{i}.message.role', '')
            content = attributes.get(f'llm.input_messages.{i}.message.content', '')
            if content:
                input_messages.append(f"[{role}]: {content}")
            i += 1

        if input_messages:
            llm_input = '\n\n'.join(input_messages)

        # Output messages - usually just one
        if 'llm.output_messages.0.message.content' in attributes:
            llm_output = attributes.get('llm.output_messages.0.message.content', '')

        # Truncate if too long (for display purposes)
        if llm_input and len(str(llm_input)) > 1000:
            llm_input = str(llm_input)[:1000] + '...'
        if llm_output and len(str(llm_output)) > 1000:
            llm_output = str(llm_output)[:1000] + '...'

        logger.debug(f"Input length: {len(str(llm_input)) if llm_input else 0}, Output length: {len(str(llm_output)) if llm_output else 0}")

        # Calculate cost using shared utility
        cost = calculate_cost(model, input_tokens, output_tokens)

        # Get current task context
        task_id = self.current_task

        # Create LLM call record
        llm_call = {
            'start_time': datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc).isoformat(),
            'end_time': datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc).isoformat(),
            'model': model,
            'agent_name': agent_name,
            'task_id': task_id,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'total_tokens': total_tokens,
            'cost': cost
        }

        self.llm_calls.append(llm_call)

        # Update task cost tracking
        if task_id and task_id in self.task_info:
            self.task_info[task_id]['input_tokens'] += input_tokens
            self.task_info[task_id]['output_tokens'] += output_tokens
            self.task_info[task_id]['cost'] += cost

        # Add span to timeline
        start_time = datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc)
        end_time = datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc)
        duration_ms = int((end_time - start_time).total_seconds() * 1000)

        # Resolve task description for this span
        task_description = None
        if task_id and task_id in self.task_info:
            task_description = self.task_info[task_id].get('description')

        self._create_step(
            step_type='llm_call',
            timestamp=start_time,
            parent_span_id=self._current_agent_span_id,
            duration_ms=duration_ms,
            agent_name=agent_name,
            model=model,
            provider=self._detect_provider(model),
            status='success',
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=round(cost, 6),
            task_id=self.current_task,
            task_description=task_description,
            description=f"LLM Call using {model}",
            input_text=llm_input,
            output_text=llm_output
        )
        logger.debug("Created span for LLM call")

    def _process_tool_call(self, span: ReadableSpan, attributes: Dict[str, Any]) -> None:
        """Extract tool call metrics from span"""

        # Log tool-related attributes for debugging
        if logger.isEnabledFor(logging.DEBUG):
            for key, value in attributes.items():
                if 'tool' in key.lower() or 'function' in key.lower():
                    logger.debug(f"Tool attr: {key}: {value}")

        # Try multiple possible attribute names for tool name
        tool_name = (
            attributes.get('tool_name') or
            attributes.get('tool.name') or
            attributes.get('crewai.tool.name') or
            attributes.get('function.name') or
            attributes.get('code.function') or
            None
        )

        # If no attribute found, try parsing from span name
        if not tool_name:
            if '.' in span.name:
                parts = span.name.split('.')
                tool_name = parts[0] if parts[0] != 'Tool' else parts[-1]
            else:
                tool_name = span.name

        logger.debug(f"Extracted tool name: {tool_name}")

        # Filter out internal tool spans
        if tool_name in INTERNAL_TOOL_NAMES or tool_name.startswith("_"):
            logger.debug(f"Skipping internal tool span: {tool_name}")
            return

        # Use current agent context
        agent_name = self.current_agent or 'unknown'

        logger.debug(f"Agent for tool: {agent_name}")

        # Extract tool input/output
        tool_input = (
            attributes.get('input.value') or
            attributes.get('tool.input') or
            attributes.get('function.arguments') or
            ''
        )
        tool_output = (
            attributes.get('output.value') or
            attributes.get('tool.output') or
            attributes.get('function.result') or
            ''
        )

        # Parse JSON input to extract just the arguments if possible
        if tool_input:
            import json
            try:
                input_data = json.loads(tool_input) if isinstance(tool_input, str) else tool_input
                if isinstance(input_data, dict) and 'calling' in input_data:
                    tool_input = input_data['calling']
            except (json.JSONDecodeError, TypeError):
                pass

        # Truncate if too long
        if tool_input and len(str(tool_input)) > 500:
            tool_input = str(tool_input)[:500] + '...'
        if tool_output and len(str(tool_output)) > 500:
            tool_output = str(tool_output)[:500] + '...'

        start_time = datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc)
        end_time = datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc)
        duration_ms = int((end_time - start_time).total_seconds() * 1000)

        # Check if tool call succeeded
        status = 'success'
        if span.status.status_code != StatusCode.OK:
            status = 'failed'

        # Also check events for exceptions
        if span.events:
            for event in span.events:
                if event.name == 'exception' or 'error' in event.name.lower():
                    status = 'failed'
                    logger.debug(f"Tool failed (detected from event: {event.name})")
                    if hasattr(event, 'attributes'):
                        error_msg = event.attributes.get('exception.message', '')
                        if error_msg:
                            tool_output = f"ERROR: {error_msg}"

        # Detect "soft errors"
        soft_error_message = None
        if tool_output and isinstance(tool_output, str):
            output_lower = tool_output.lower()
            if 'i encountered an error' in output_lower or 'this was the error:' in output_lower:
                status = 'failed'
                if 'this was the error:' in output_lower:
                    try:
                        start_idx = tool_output.lower().index('this was the error:') + len('this was the error:')
                        end_idx = tool_output.find('.', start_idx + 1)
                        if end_idx > start_idx:
                            soft_error_message = tool_output[start_idx:end_idx].strip()
                    except (ValueError, IndexError):
                        soft_error_message = "Tool execution failed"
                else:
                    soft_error_message = "Tool execution failed"

                # Create error entry for soft errors
                error_info = {
                    'timestamp': end_time.isoformat(),
                    'error_type': 'TOOL_ERROR',
                    'message': soft_error_message or "Tool execution failed",
                    'agent_name': agent_name,
                    'tool_name': tool_name
                }
                if self.current_task:
                    error_info['task_id'] = self.current_task
                self.errors.append(error_info)
                logger.debug(f"Soft error captured: {soft_error_message}")

        tool_call = {
            'tool_name': tool_name,
            'agent_name': agent_name,
            'execution_time_ms': duration_ms,
            'status': status,
            'start_time': start_time.isoformat(),
            'end_time': end_time.isoformat(),
            'input': str(tool_input) if tool_input else None,
            'output': str(tool_output) if tool_output else None
        }

        # Check if we should merge with an existing tool call
        merged = False
        for existing in self.tool_calls:
            if existing['tool_name'] == tool_name and existing['agent_name'] == agent_name:
                existing_start = datetime.fromisoformat(existing['start_time'])
                if abs((start_time - existing_start).total_seconds()) < 0.2:
                    if tool_input and not existing.get('input'):
                        existing['input'] = str(tool_input)
                    if tool_output and not existing.get('output'):
                        existing['output'] = str(tool_output)
                    merged = True
                    logger.debug("Merged tool data into existing call")
                    break

        if not merged:
            self.tool_calls.append(tool_call)

            # Resolve task description for this span
            task_description = None
            if self.current_task and self.current_task in self.task_info:
                task_description = self.task_info[self.current_task].get('description')

            self._create_step(
                step_type='tool_call',
                timestamp=start_time,
                parent_span_id=self._current_agent_span_id,
                duration_ms=duration_ms,
                agent_name=agent_name,
                tool_name=tool_name,
                status=status,
                task_id=self.current_task,
                task_description=task_description,
                description=f"Tool: {tool_name}",
                input_text=tool_input,
                output_text=tool_output
            )
            logger.debug(f"Created span for tool call (status: {status})")

    def _process_agent_execution(self, span: ReadableSpan, attributes: Dict[str, Any]) -> None:
        """Track agent execution timing"""
        agent_name = (
            attributes.get('graph.node.id') or
            attributes.get('agent_role') or
            attributes.get('crewai.agent.name') or
            attributes.get('crewai.agent.role') or
            f"Agent_{len(self.agents_executed) + 1}"
        )

        logger.debug(f"Agent execution: {agent_name}")

        # Set current agent context
        self.current_agent = agent_name

        if agent_name not in self.agents_executed:
            self.agents_executed.append(agent_name)

        if agent_name not in self.agent_timings:
            self.agent_timings[agent_name] = {
                'start_time': datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc),
                'end_time': datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc),
                'duration_ms': 0
            }
        else:
            self.agent_timings[agent_name]['end_time'] = datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc)

        # Calculate duration
        start = self.agent_timings[agent_name]['start_time']
        end = self.agent_timings[agent_name]['end_time']
        self.agent_timings[agent_name]['duration_ms'] = int((end - start).total_seconds() * 1000)

        # Add step to timeline and track span_id for parent-child linking
        if agent_name not in [s.get('agent_name') for s in self.steps if s.get('type') == 'agent_start']:
            start_time = datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc)
            agent_span_id = self._create_step(
                step_type='agent_start',
                timestamp=start_time,
                agent_name=agent_name,
                description=f"Agent started: {agent_name}"
            )
            self._current_agent_span_id = agent_span_id
            logger.debug("Created span for agent start")
        else:
            # Agent already has a span — find its span_id
            for s in self.steps:
                if s.get('type') == 'agent_start' and s.get('agent_name') == agent_name:
                    self._current_agent_span_id = s.get('span_id')
                    break

    def _classify_error_type(self, raw_type: str, error_message: str, tool_name: str = None) -> str:
        """Classify raw error type into standardized categories"""
        raw_lower = (raw_type + ' ' + error_message).lower()

        if tool_name:
            return "TOOL_ERROR"
        if 'timeout' in raw_lower or 'timed out' in raw_lower:
            return "TIMEOUT"
        if any(kw in raw_lower for kw in ['rate', 'api', 'openai', 'anthropic', 'llm', 'model', 'token limit', 'context length']):
            return "LLM_ERROR"
        if any(kw in raw_lower for kw in ['valid', 'schema', 'format', 'parse', 'type error']):
            return "VALIDATION"
        return "UNKNOWN"

    def _process_error(self, span: ReadableSpan, attributes: Dict[str, Any]) -> None:
        """Extract error information from failed span"""
        raw_error_type = attributes.get('exception.type', 'UnknownError')
        error_message = span.status.description or attributes.get('exception.message', 'No error message')
        stack_trace = attributes.get('exception.stacktrace', '')

        agent_name = self.current_agent or (
            attributes.get('agent_role') or
            attributes.get('graph.node.id') or
            'unknown'
        )

        tool_name = attributes.get('crewai.tool.name')
        error_type = self._classify_error_type(raw_error_type, error_message, tool_name)

        error_info = {
            'timestamp': datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc).isoformat(),
            'error_type': error_type,
            'message': error_message,
            'agent_name': agent_name
        }

        if tool_name:
            error_info['tool_name'] = tool_name

        if self.current_task:
            error_info['task_id'] = self.current_task

        if stack_trace:
            error_info['stack_trace'] = stack_trace[:500] if len(stack_trace) > 500 else stack_trace

        self.errors.append(error_info)
        logger.debug(f"Error captured: {error_type} - {error_message}")

        # Create error step in timeline
        timestamp = datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc)
        description = f"Error: {error_type}"
        if tool_name:
            description = f"Tool Error: {tool_name} - {error_type}"

        self._create_step(
            step_type='error',
            timestamp=timestamp,
            parent_span_id=self._current_agent_span_id,
            agent_name=agent_name,
            error_type=error_type,
            error_message=error_message,
            tool_name=tool_name,
            task_id=self.current_task,
            stack_trace=stack_trace[:500] if stack_trace and len(stack_trace) > 500 else stack_trace,
            description=description
        )
        logger.debug("Created error span in timeline")

    def _deduplicate_errors(self, errors: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove duplicate errors within 100ms window."""
        if not errors:
            return errors

        seen = {}
        unique_errors = []

        for error in sorted(errors, key=lambda e: e['timestamp']):
            key = (error['message'], error.get('tool_name'))

            timestamp_str = error['timestamp']
            if timestamp_str.endswith('Z'):
                timestamp_str = timestamp_str[:-1] + '+00:00'
            try:
                timestamp = datetime.fromisoformat(timestamp_str)
            except ValueError:
                unique_errors.append(error)
                continue

            if key in seen:
                last_time = seen[key]
                if abs((timestamp - last_time).total_seconds()) < 0.1:
                    continue

            unique_errors.append(error)
            seen[key] = timestamp

        return unique_errors

    def get_metrics(self) -> Dict[str, Any]:
        """Get aggregate metrics for trace completion.

        Breakdowns (agents, tools, errors) are computed server-side from spans.
        This only returns the totals needed for the PATCH /api/traces/:id call.
        """
        if not self.start_time or not self.end_time:
            return {
                'duration_ms': 0,
                'end_time': datetime.now(timezone.utc).isoformat(),
                'total_cost': 0.0,
                'total_input_tokens': 0,
                'total_output_tokens': 0,
            }

        duration_ms = int((self.end_time - self.start_time).total_seconds() * 1000)

        # Calculate totals from accumulated LLM calls
        total_input_tokens = sum(c['input_tokens'] for c in self.llm_calls)
        total_output_tokens = sum(c['output_tokens'] for c in self.llm_calls)
        total_cost = sum(c['cost'] for c in self.llm_calls)

        logger.debug(
            f"Final metrics: agents={self.agents_executed}, "
            f"cost=${total_cost:.6f}, llm_calls={len(self.llm_calls)}, "
            f"tool_calls={len(self.tool_calls)}, spans={len(self.steps)}"
        )

        return {
            'duration_ms': duration_ms,
            'end_time': self.end_time.isoformat(),
            'total_cost': round(total_cost, 6),
            'total_input_tokens': total_input_tokens,
            'total_output_tokens': total_output_tokens,
        }

    def shutdown(self) -> None:
        """Shutdown the exporter"""
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush any buffered spans"""
        return True


class CrewTracker:
    """
    Context manager for tracking CrewAI execution using OpenTelemetry.

    Usage:
        with tracker.track(crew, "workflow-name"):
            result = crew.kickoff()
    """

    def __init__(self, client, crew, name: str):
        """
        Initialize the crew tracker.

        Args:
            client: Visibe client instance
            crew: CrewAI Crew instance
            name: Name for this execution trace
        """
        self.client = client
        self.crew = crew
        self.name = name
        self.trace_id = str(uuid.uuid4())
        self.exporter = CrewAISpanExporter()
        self.exporter.populate_task_descriptions(crew)
        self._span_processor = None  # Set by track() after _add_span_processor

    def __enter__(self):
        """Enter — create trace header on backend and wire span streaming."""
        global _active_crew_trace

        self.client.create_trace({
            'trace_id': self.trace_id,
            'name': self.name,
            'framework': 'crewai',
            'started_at': datetime.now(timezone.utc).isoformat(),
        })

        # Wire real-time span streaming
        self.exporter._span_sender = lambda span: self.client.queue_span(
            self.trace_id, span
        )

        # Signal to the OpenAI patcher that a CrewAI trace is active
        _active_crew_trace = True

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit — flush OTEL spans, complete trace with summary."""
        global _active_crew_trace
        _active_crew_trace = False

        # Force flush all pending OTEL spans so exporter.export() is called
        tracer_provider = trace.get_tracer_provider()
        if hasattr(tracer_provider, 'force_flush'):
            try:
                tracer_provider.force_flush(timeout_millis=5000)
            except Exception as e:
                logger.warning(f"Force flush failed: {e}")

        # Collect aggregated metrics for trace completion
        metrics = self.exporter.get_metrics()

        # Extract prompt from first task's description
        prompt = None
        if hasattr(self.crew, 'tasks') and self.crew.tasks:
            first_task = self.crew.tasks[0]
            if hasattr(first_task, 'description'):
                prompt = first_task.description

        # Determine status
        has_errors = len(self.exporter.errors) > 0
        status = 'failed' if (exc_type is not None or has_errors) else 'completed'

        total_input_tokens = metrics['total_input_tokens']
        total_output_tokens = metrics['total_output_tokens']
        total_cost = metrics['total_cost']
        total_tokens = total_input_tokens + total_output_tokens
        model = self.exporter.llm_calls[0]['model'] if self.exporter.llm_calls else None

        # Complete trace with summary (spans already streamed)
        success = self.client.complete_trace(self.trace_id, {
            'status': status,
            'ended_at': metrics['end_time'],
            'duration_ms': metrics['duration_ms'],
            'prompt': prompt or "No prompt captured",
            'model': model,
            'total_cost': total_cost,
            'total_tokens': total_tokens,
            'total_input_tokens': total_input_tokens,
            'total_output_tokens': total_output_tokens,
        })

        # Print summary to stdout
        duration_s = metrics['duration_ms'] / 1000
        llm_calls = len(self.exporter.llm_calls)
        tool_calls_count = len(self.exporter.tool_calls)

        summary = TraceSummary(
            name=self.name,
            status=status,
            llm_calls=llm_calls,
            tool_calls=tool_calls_count,
            total_tokens=total_tokens,
            total_cost=total_cost,
            duration_s=duration_s,
            agents=self.exporter.agents_executed,
            sent=success,
        )
        print(f"[Visibe] {summary}")

        # Clean up span processor to prevent accumulation
        if self._span_processor is not None:
            try:
                self._span_processor.shutdown()
            except Exception:
                pass

        return False
