"""LM instrumentation for nanomon."""

from __future__ import annotations

import logging
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from dspy.clients.base_lm import BaseLM

from nanomon.usage_extractors import ExtractorChain
from nanomon.utils import Clock, SystemClock, ns_to_ms

if TYPE_CHECKING:
    from nanomon.sinks import Sink
    from nanomon.types import RunState

logger = logging.getLogger(__name__)


@dataclass
class ActiveContext:
    """Holds context info and the current run being accumulated.

    This is set by NanomonRunContext.run() and the NanomonCallback manages
    the current_run_state lifecycle based on DSPy module boundaries.
    """

    context_id: str
    sink: Sink
    clock: Clock
    extractor_chain: ExtractorChain
    tags: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    provider: str | None = None
    current_run_state: RunState | None = None  # Set by callback when module starts

    @property
    def current_run_id(self) -> str | None:
        """Get the ID of the currently active run (for tool tracking)."""
        return self.current_run_state.id if self.current_run_state else None


# Context variable holding the active context (if any)
_active_context: ContextVar[ActiveContext | None] = ContextVar("active_context", default=None)


def get_active_context() -> ActiveContext | None:
    """Get the currently active context, if any."""
    return _active_context.get()


def set_active_context(ctx: ActiveContext | None) -> None:
    """Set the active context."""
    _active_context.set(ctx)


def get_current_run_id() -> str | None:
    """Get the ID of the currently active run (for tool tracking).

    Returns None if no context is active or no run is in progress.
    """
    ctx = _active_context.get()
    return ctx.current_run_id if ctx else None


# Keep old functions for backwards compatibility
def get_current_run() -> Any:
    """Deprecated: Get the current run state.

    For tool tracking, use get_current_run_id() instead.
    Returns a minimal object with just the id attribute.
    """
    ctx = _active_context.get()
    if ctx and ctx.current_run_id:
        # Return a minimal object that has .id for backwards compat
        class _RunRef:
            def __init__(self, run_id: str):
                self.id = run_id

        return _RunRef(ctx.current_run_id)
    return None


def set_current_run(run: Any) -> None:
    """Deprecated: No longer used. Context is set via set_active_context."""
    pass


class InstrumentedLM(BaseLM):
    """Wrapper around a DSPy LM that intercepts calls for metrics collection.

    This wrapper inherits from dspy.BaseLM to pass DSPy's type checks while
    capturing metrics. LLM calls accumulate stats to the current RunState
    which is managed by the NanomonCallback based on DSPy module boundaries.
    """

    def __init__(
        self,
        lm: BaseLM,
        extractor_chain: ExtractorChain,
        clock: Clock | None = None,
    ) -> None:
        """Initialize the instrumented LM.

        Args:
            lm: The original DSPy LM to wrap
            extractor_chain: Chain of extractors for usage/model extraction
            clock: Clock for timing (defaults to SystemClock)
        """
        # Initialize BaseLM with the wrapped LM's attributes
        super().__init__(
            model=lm.model,
            model_type=getattr(lm, "model_type", "chat"),
            temperature=lm.kwargs.get("temperature", 0.0),
            max_tokens=lm.kwargs.get("max_tokens", 1000),
            cache=getattr(lm, "cache", True),
            **{k: v for k, v in lm.kwargs.items() if k not in ("temperature", "max_tokens")},
        )

        self._wrapped_lm = lm
        self._extractor_chain = extractor_chain
        self._clock = clock or SystemClock()

    def forward(
        self, prompt: str | None = None, messages: list[dict[str, Any]] | None = None, **kwargs: Any
    ) -> Any:
        """Forward pass that accumulates metrics to the current run.

        Args:
            prompt: The prompt string
            messages: The messages list
            **kwargs: Additional arguments to pass to the LM

        Returns:
            The original LM response (unmodified)
        """
        ctx = get_active_context()

        # If no context or no active run, just call LLM without tracking
        if ctx is None or ctx.current_run_state is None:
            return self._wrapped_lm.forward(prompt=prompt, messages=messages, **kwargs)

        start_ns = self._clock.perf_counter_ns()
        response = None

        try:
            response = self._wrapped_lm.forward(prompt=prompt, messages=messages, **kwargs)
            return response
        finally:
            # Only accumulate if we got a response
            if response is not None:
                end_ns = self._clock.perf_counter_ns()
                latency_ms = ns_to_ms(end_ns - start_ns)

                # Extract metrics
                input_tokens, output_tokens = ctx.extractor_chain.extract_usage(response)
                model_name = ctx.extractor_chain.extract_model_name(
                    response, self._wrapped_lm, kwargs
                )

                # ACCUMULATE to current RunState (no API call here!)
                ctx.current_run_state.add_call(
                    model_name, input_tokens, output_tokens, latency_ms
                )

                logger.debug(
                    f"Accumulated LLM call: model={model_name}, "
                    f"tokens={input_tokens}+{output_tokens}, "
                    f"latency={latency_ms}ms"
                )

    async def aforward(
        self, prompt: str | None = None, messages: list[dict[str, Any]] | None = None, **kwargs: Any
    ) -> Any:
        """Async forward pass that accumulates metrics to the current run.

        Args:
            prompt: The prompt string
            messages: The messages list
            **kwargs: Additional arguments to pass to the LM

        Returns:
            The original LM response (unmodified)
        """
        ctx = get_active_context()

        # If no context or no active run, just call LLM without tracking
        if ctx is None or ctx.current_run_state is None:
            return await self._wrapped_lm.aforward(prompt=prompt, messages=messages, **kwargs)

        start_ns = self._clock.perf_counter_ns()
        response = None

        try:
            response = await self._wrapped_lm.aforward(prompt=prompt, messages=messages, **kwargs)
            return response
        finally:
            # Only accumulate if we got a response
            if response is not None:
                end_ns = self._clock.perf_counter_ns()
                latency_ms = ns_to_ms(end_ns - start_ns)

                # Extract metrics
                input_tokens, output_tokens = ctx.extractor_chain.extract_usage(response)
                model_name = ctx.extractor_chain.extract_model_name(
                    response, self._wrapped_lm, kwargs
                )

                # ACCUMULATE to current RunState (no API call here!)
                ctx.current_run_state.add_call(
                    model_name, input_tokens, output_tokens, latency_ms
                )

                logger.debug(
                    f"Accumulated LLM call: model={model_name}, "
                    f"tokens={input_tokens}+{output_tokens}, "
                    f"latency={latency_ms}ms"
                )

    def __getattr__(self, name: str) -> Any:
        """Proxy attribute access to the underlying LM for attributes not on this class."""
        # Avoid infinite recursion for _wrapped_lm itself
        if name == "_wrapped_lm":
            raise AttributeError(name)
        return getattr(self._wrapped_lm, name)

    def __repr__(self) -> str:
        return f"InstrumentedLM({self._wrapped_lm!r})"
