# Copyright (c) Metis. All rights reserved.

"""Span processors for MantisDK tracing.

This module provides SpanProcessor implementations that inject context
into spans automatically. The TracingContextSpanProcessor reads from
context variables and injects attributes into every span.
"""

from __future__ import annotations

import json
import logging
from typing import Optional

from opentelemetry.context import Context
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
from opentelemetry.trace import Span, StatusCode

from . import semconv

logger = logging.getLogger(__name__)


class TracingContextSpanProcessor(SpanProcessor):
    """Injects all tracing context into spans automatically.

    This processor runs on every span creation and injects context from:

    1. Run module (run_id, session_id from current run/rollout)
    2. Context module (call_type, tags, environment)
    3. Experiments module (experiment_id, variant_key, user_id)

    Attributes injected:
        - insight.run.id: Current run ID (from mantisdk.run)
        - insight.run.name: Current run name (from mantisdk.run)
        - insight.session.id: Current rollout session ID (from mantisdk.run)
        - mantis.call_type: Call type (e.g., "agent-call", "judge-call")
        - insight.trace.tags: JSON-encoded list of tags
        - insight.environment: Environment name
        - msk.experiment.id: A/B experiment key (from mantisdk.experiments)
        - msk.variant.key: Assigned variant key (from mantisdk.experiments)
        - msk.experiment.user_id: User ID for experiment assignment

    Example::

        from opentelemetry.sdk.trace import TracerProvider
        from mantisdk.tracing.processors import TracingContextSpanProcessor

        provider = TracerProvider()
        provider.add_span_processor(TracingContextSpanProcessor())

        # All spans will now have context injected automatically
    """

    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
        """Called when a span is started.

        Injects context attributes from run module, context module, and experiments module.
        Failures are logged at debug level and do not raise exceptions.
        """
        if not span.is_recording():
            return

        try:
            self._inject_run_context(span)
        except ImportError:
            # run module not available, skip run context injection
            pass
        except Exception as e:
            logger.debug(f"Could not inject run context: {e}")

        try:
            self._inject_tracing_context(span)
        except ImportError:
            # context module not available, skip tracing context injection
            pass
        except Exception as e:
            logger.debug(f"Could not inject tracing context: {e}")

        try:
            self._inject_experiment_context(span)
        except ImportError:
            # experiments module not available, skip experiment context injection
            pass
        except Exception as e:
            logger.debug(f"Could not inject experiment context: {e}")

    def _inject_run_context(self, span: Span) -> None:
        """Inject run and rollout context into span."""
        from mantisdk.run import get_current_rollout, get_current_run

        # Run context
        current_run = get_current_run()
        if current_run is not None:
            if current_run.id:
                span.set_attribute(semconv.RUN_ID, current_run.id)
            if current_run.name:
                span.set_attribute(semconv.RUN_NAME, current_run.name)

        # Rollout context -> session_id + trace_id capture
        current_rollout = get_current_rollout()
        if current_rollout is not None and current_rollout.session_id:
            span.set_attribute(semconv.SESSION_ID, current_rollout.session_id)

            # Capture the trace_id from this span onto the Rollout so that
            # rollout.score() can link scores to the correct trace even
            # after instrumented spans have ended.
            span_context = span.get_span_context()
            if span_context and span_context.is_valid:
                trace_id_hex = format(span_context.trace_id, "032x")
                current_rollout._capture_trace_id(trace_id_hex)
        else:
            # Fallback: proxy-specific session (set by RolloutAttemptMiddleware)
            from .context import get_proxy_session_id

            proxy_session = get_proxy_session_id()
            if proxy_session:
                span.set_attribute(semconv.SESSION_ID, proxy_session)

    def _inject_tracing_context(self, span: Span) -> None:
        """Inject call_type, tags, and environment context into span."""
        from .context import (
            get_current_call_type,
            get_current_environment,
            get_current_tags,
        )

        # Call type (from @gepa.agent, @gepa.judge, etc.)
        call_type = get_current_call_type()
        if call_type:
            span.set_attribute("mantis.call_type", call_type)

        # Tags
        tags = get_current_tags()
        if tags:
            span.set_attribute(semconv.TRACE_TAGS, json.dumps(tags))

        # Environment
        environment = get_current_environment()
        if environment:
            span.set_attribute(semconv.ENVIRONMENT, environment)

    def _inject_experiment_context(self, span: Span) -> None:
        """Inject A/B experiment context into span.

        When get_variant() is called, it sets the current experiment context.
        This method reads that context and injects the experiment_id and
        variant_key as span attributes so they flow to ClickHouse.

        Attribute names must match attributeAliases.ts in the Insight backend:
        - msk.experiment.id -> ab_experiment.id
        - msk.variant.key -> ab_experiment.variant_key
        - msk.experiment.user_id -> ab_experiment.user_id
        """
        from mantisdk.experiments.client import get_current_experiment

        experiment_ctx = get_current_experiment()
        if experiment_ctx is None:
            return

        experiment_id = experiment_ctx.get("experiment_id")
        variant_key = experiment_ctx.get("variant_key")
        user_id = experiment_ctx.get("user_id")

        if experiment_id:
            span.set_attribute("msk.experiment.id", experiment_id)
        if variant_key:
            span.set_attribute("msk.variant.key", variant_key)
        if user_id:
            span.set_attribute("msk.experiment.user_id", user_id)

    def on_end(self, span: ReadableSpan) -> None:
        """Called when a span is ended. No-op for this processor."""
        pass

    def shutdown(self) -> None:
        """Called when the TracerProvider is shut down. No-op for this processor."""
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush any buffered data. Always returns True for this processor."""
        return True


class LoggingSpanProcessor(SpanProcessor):
    """Logs span lifecycle events for debug visibility.

    When added to a TracerProvider (via ``tracing.init(debug=True)``), this
    processor prints span start/end to the ``mantisdk.tracing`` logger so
    users can see what the tracing pipeline is doing.

    Auto-instrumented spans (OpenAI, Anthropic, etc.) are identified by
    their instrumentation scope name and labelled accordingly.
    """

    _TRACER_NAME = "mantisdk.tracing"

    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
        if not span.is_recording():
            return
        scope = getattr(span, "instrumentation_scope", None)
        scope_name = getattr(scope, "name", None) if scope else None
        if scope_name and scope_name != self._TRACER_NAME:
            source = f" (auto: {scope_name.rsplit('.', 1)[-1]})"
        else:
            source = ""
        logger.debug('Span started: "%s"%s', span.name, source)

    def on_end(self, span: ReadableSpan) -> None:
        duration_s = (span.end_time - span.start_time) / 1e9 if span.end_time and span.start_time else 0
        ctx = span.context
        trace_id = format(ctx.trace_id, "032x") if ctx else ""
        status = span.status
        if status and status.status_code == StatusCode.ERROR:
            suffix = f" [ERROR: {status.description or ''}]"
        else:
            suffix = ""
        logger.debug('Span ended:   "%s" (%.2fs) trace_id=%s%s', span.name, duration_s, trace_id, suffix)

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return True


class AnnotationFilterSpanProcessor(SpanProcessor):
    """Wraps another SpanProcessor, filtering out mantisdk.annotation spans.

    Reward annotations are internal to mantisdk (used by the store's
    ``find_final_reward()``). They are also sent as proper Langfuse Scores
    via the API by ``InsightRunListener``. Exporting them as OTLP traces
    creates noise (null input, undefined output) without adding value.

    This processor drops ``mantisdk.annotation`` spans before they reach the
    wrapped exporter while letting all other spans pass through.

    Example::

        from opentelemetry.sdk.trace.export import BatchSpanProcessor
        from mantisdk.tracing.processors import AnnotationFilterSpanProcessor

        inner = BatchSpanProcessor(exporter)
        filtered = AnnotationFilterSpanProcessor(inner)
        provider.add_span_processor(filtered)
    """

    _FILTERED_NAME = "mantisdk.annotation"

    def __init__(self, inner: SpanProcessor) -> None:
        self._inner = inner

    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
        self._inner.on_start(span, parent_context)

    def on_end(self, span: ReadableSpan) -> None:
        if span.name == self._FILTERED_NAME:
            return  # Drop reward annotation spans from OTLP export
        self._inner.on_end(span)

    def shutdown(self) -> None:
        self._inner.shutdown()

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return self._inner.force_flush(timeout_millis)
