"""
DSPy callback for Risicare SDK.

Implements dspy.BaseCallback handlers for module, LM, and tool events.
Uses call_id-based span tracking with memory guard (same pattern as
LangChain's _run_spans).

All attribute access is defensive (getattr with defaults) since DSPy's
callback handler signatures are under-documented and may vary between
versions 2.x and 3.x.
"""

from __future__ import annotations

import logging
import threading
import time
from typing import Any, Dict, Optional

import wrapt

from risicare.integrations._base import (
    create_framework_attributes,
    get_framework_version,
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
)
from risicare.integrations._dedup import (
    suppress_provider_instrumentation,
)

logger = logging.getLogger(__name__)

# Memory guard constants (same as LangChain)
_MAX_OPEN_SPANS = 10_000
_SPAN_TTL_SECONDS = 300  # 5 minutes


class _SpanEntry:
    """Tracks a span and its metadata for correlation."""

    __slots__ = ("span", "suppression_cm", "start_time", "span_kind")

    def __init__(
        self,
        span: Any,
        suppression_cm: Any,
        start_time: float,
        span_kind: str,
    ) -> None:
        self.span = span
        self.suppression_cm = suppression_cm
        self.start_time = start_time
        self.span_kind = span_kind


class RisicareDSPyCallback:
    """
    DSPy callback that creates Risicare spans.

    We don't inherit from dspy.BaseCallback to avoid requiring dspy
    as an import-time dependency. DSPy's callback system uses duck
    typing — it checks for method existence.
    """

    def __init__(self) -> None:
        self._version = get_framework_version("dspy")
        self._open_spans: Dict[str, _SpanEntry] = {}
        self._lock = threading.Lock()

    def _evict_stale_entries(self) -> None:
        """Evict stale entries from _open_spans to prevent memory leaks."""
        if len(self._open_spans) < _MAX_OPEN_SPANS:
            return

        now = time.perf_counter()
        stale_keys = [
            k for k, v in self._open_spans.items()
            if (now - v.start_time) > _SPAN_TTL_SECONDS
        ]

        if stale_keys:
            for key in stale_keys:
                entry = self._open_spans.pop(key, None)
                if entry:
                    self._cleanup_entry(entry)
        else:
            # No stale entries — FIFO eviction of oldest 10%
            evict_count = max(1, len(self._open_spans) // 10)
            sorted_keys = sorted(
                self._open_spans.keys(),
                key=lambda k: self._open_spans[k].start_time,
            )
            for key in sorted_keys[:evict_count]:
                entry = self._open_spans.pop(key, None)
                if entry:
                    self._cleanup_entry(entry)

    def _start_span(
        self,
        call_id: str,
        instance: Any,
        inputs: Any,
        span_kind: str,
        span_name: str,
        suppress_providers: bool = False,
    ) -> None:
        """Create and track a new span."""
        if not is_tracing_enabled():
            return

        tracer = get_tracer()
        if tracer is None:
            return

        try:
            from risicare_core import SpanKind as RisicareSpanKind
        except ImportError:
            return

        with self._lock:
            self._evict_stale_entries()

        attrs = create_framework_attributes("dspy", self._version)
        attrs["framework.span_kind"] = span_kind

        # Defensive attribute extraction
        if span_kind == "module":
            module_type = type(instance).__name__ if instance else "unknown"
            attrs["framework.dspy.module_type"] = module_type
            signature = getattr(instance, "signature", None)
            if signature:
                attrs["framework.dspy.signature"] = str(signature)[:200]
            kind = RisicareSpanKind.INTERNAL
        elif span_kind == "lm":
            model = getattr(instance, "model", None) or getattr(instance, "model_name", None)
            if model:
                attrs["gen_ai.request.model"] = str(model)
                # Extract provider from model prefix (e.g., "openai/gpt-4o" → "openai")
                if "/" in str(model):
                    attrs["gen_ai.system"] = str(model).split("/")[0]
            kind = RisicareSpanKind.LLM_CALL
        elif span_kind == "tool":
            tool_name = getattr(instance, "name", None) or type(instance).__name__
            attrs["framework.dspy.tool_name"] = str(tool_name)
            kind = RisicareSpanKind.TOOL_CALL
        else:
            kind = RisicareSpanKind.INTERNAL

        if should_trace_content() and inputs:
            try:
                inputs_str = str(inputs)
                attrs["gen_ai.request.messages"] = truncate_content(inputs_str)
            except Exception:
                pass

        span = tracer.start_span_no_context(
            name=span_name,
            kind=kind,
            attributes=attrs,
        )

        suppression_cm = None
        if suppress_providers:
            suppression_cm = suppress_provider_instrumentation()
            suppression_cm.__enter__()

        try:
            with self._lock:
                self._open_spans[str(call_id)] = _SpanEntry(
                    span=span,
                    suppression_cm=suppression_cm,
                    start_time=time.perf_counter(),
                    span_kind=span_kind,
                )
        except BaseException:
            if suppression_cm is not None:
                suppression_cm.__exit__(None, None, None)
            raise

    def _end_span(
        self,
        call_id: str,
        outputs: Any = None,
        exception: Any = None,
    ) -> None:
        """End a tracked span."""
        with self._lock:
            entry = self._open_spans.pop(str(call_id), None)
        if entry is None:
            return

        try:
            span = entry.span
            latency_ms = (time.perf_counter() - entry.start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if exception is not None:
                if isinstance(exception, BaseException):
                    record_error(span, exception)
                else:
                    span.set_attribute("error", True)
                    span.set_attribute("error.message", str(exception)[:500])
            elif outputs is not None:
                if should_trace_content():
                    try:
                        outputs_str = str(outputs)
                        safe_set_attribute(
                            span,
                            "gen_ai.response.content",
                            truncate_content(outputs_str),
                        )
                    except Exception:
                        pass

                # Try to extract token usage from LM outputs
                if entry.span_kind == "lm":
                    try:
                        usage = getattr(outputs, "usage", None)
                        if usage:
                            safe_set_attribute(span, "gen_ai.usage.prompt_tokens", getattr(usage, "prompt_tokens", None))
                            safe_set_attribute(span, "gen_ai.usage.completion_tokens", getattr(usage, "completion_tokens", None))
                    except Exception:
                        pass

            span.end()
        except Exception as e:
            logger.debug(f"Error ending DSPy span: {e}")
        finally:
            self._cleanup_entry(entry)

    def _cleanup_entry(self, entry: _SpanEntry) -> None:
        """Clean up suppression context."""
        if entry.suppression_cm is not None:
            try:
                entry.suppression_cm.__exit__(None, None, None)
            except Exception:
                pass

    # =====================================================================
    # DSPy callback handlers
    # =====================================================================

    def on_module_start(self, call_id: Any, instance: Any, inputs: Any) -> None:
        """Called when a DSPy Module is invoked (e.g., ChainOfThought)."""
        module_type = type(instance).__name__ if instance else "module"
        self._start_span(
            call_id=str(call_id),
            instance=instance,
            inputs=inputs,
            span_kind="module",
            span_name=f"dspy.module/{module_type}",
            suppress_providers=False,  # Module spans don't suppress — LM spans do
        )

    def on_module_end(self, call_id: Any, outputs: Any, exception: Any = None) -> None:
        """Called when a DSPy Module invocation completes."""
        self._end_span(str(call_id), outputs=outputs, exception=exception)

    def on_lm_start(self, call_id: Any, instance: Any, inputs: Any) -> None:
        """Called when a DSPy LM call starts."""
        model = getattr(instance, "model", None) or "unknown"
        self._start_span(
            call_id=str(call_id),
            instance=instance,
            inputs=inputs,
            span_kind="lm",
            span_name=f"dspy.lm/{model}",
            suppress_providers=True,  # Suppress underlying SDK spans
        )

    def on_lm_end(self, call_id: Any, outputs: Any, exception: Any = None) -> None:
        """Called when a DSPy LM call completes."""
        self._end_span(str(call_id), outputs=outputs, exception=exception)

    def on_tool_start(self, call_id: Any, instance: Any, inputs: Any) -> None:
        """Called when a DSPy Tool is invoked."""
        tool_name = getattr(instance, "name", None) or type(instance).__name__ if instance else "tool"
        self._start_span(
            call_id=str(call_id),
            instance=instance,
            inputs=inputs,
            span_kind="tool",
            span_name=f"dspy.tool/{tool_name}",
            suppress_providers=False,  # Tools don't call LLM SDKs
        )

    def on_tool_end(self, call_id: Any, outputs: Any, exception: Any = None) -> None:
        """Called when a DSPy Tool invocation completes."""
        self._end_span(str(call_id), outputs=outputs, exception=exception)


# Module-level singleton
_callback_instance: Optional[RisicareDSPyCallback] = None


def _wrap_configure(
    wrapped: Any,
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> Any:
    """Wrapt wrapper for dspy.configure — injects Risicare callback."""
    callbacks = kwargs.get("callbacks", None)
    if callbacks is None:
        kwargs["callbacks"] = [_callback_instance]
    elif isinstance(callbacks, list):
        if _callback_instance not in callbacks:
            callbacks.append(_callback_instance)
    return wrapped(*args, **kwargs)


def register_callback(module: Any) -> None:
    """
    Register the Risicare callback with DSPy.

    Uses wrapt to patch dspy.configure, auto-injecting our callback
    when the user calls dspy.configure().
    """
    global _callback_instance

    if _callback_instance is not None:
        return

    _callback_instance = RisicareDSPyCallback()

    # Wrap dspy.configure using wrapt (not raw monkey-patch)
    if getattr(module, "configure", None) is None:
        logger.debug("dspy.configure not found — cannot register callback")
        return

    try:
        wrapt.wrap_function_wrapper(module, "configure", _wrap_configure)
    except Exception as e:
        logger.debug(f"Failed to patch dspy.configure: {e}")
        return

    # Also try to register immediately if settings are already configured
    try:
        settings = getattr(module, "settings", None)
        if settings is not None:
            existing_callbacks = getattr(settings, "callbacks", None)
            if existing_callbacks is None:
                pass  # Will be set on next configure() call
            elif isinstance(existing_callbacks, list):
                if _callback_instance not in existing_callbacks:
                    existing_callbacks.append(_callback_instance)
    except Exception:
        pass

    logger.debug("Registered Risicare DSPy callback")
