"""DSPy callback for cost logging run boundaries."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

import dspy
from dspy.utils.callback import ACTIVE_CALL_ID, BaseCallback

from nanomon.instrumentation import set_active_context
from nanomon.types import RunState
from nanomon.utils import generate_uuid

if TYPE_CHECKING:
    from nanomon.instrumentation import ActiveContext

logger = logging.getLogger(__name__)


def _serialize_value(value: Any) -> Any:
    """Serialize a value to JSON-compatible format."""
    if isinstance(value, (str, int, float, bool, type(None))):
        return value
    if isinstance(value, (list, tuple)):
        return [_serialize_value(v) for v in value]
    if isinstance(value, dict):
        return {k: _serialize_value(v) for k, v in value.items()}
    # Handle dspy.Prediction and similar objects
    if hasattr(value, "toDict"):
        return value.toDict()
    if hasattr(value, "__dict__"):
        return {k: _serialize_value(v) for k, v in value.__dict__.items() if not k.startswith("_")}
    return str(value)  # Fallback to string representation


def _serialize_inputs(inputs: dict[str, Any]) -> dict[str, Any]:
    """Serialize inputs dictionary to JSON-compatible format."""
    return {k: _serialize_value(v) for k, v in inputs.items()}


def _serialize_outputs(outputs: Any) -> dict[str, Any] | None:
    """Serialize DSPy outputs to JSON-compatible format."""
    if outputs is None:
        return None
    if hasattr(outputs, "toDict"):
        return outputs.toDict()
    if isinstance(outputs, dict):
        return {k: _serialize_value(v) for k, v in outputs.items()}
    if hasattr(outputs, "__dict__"):
        return {k: _serialize_value(v) for k, v in outputs.__dict__.items() if not k.startswith("_")}
    return {"value": _serialize_value(outputs)}


class NanomonCallback(BaseCallback):
    """DSPy callback that creates runs per top-level module invocation.

    This callback hooks into DSPy's module execution lifecycle:
    - on_module_start: Creates a new run when a TOP-LEVEL ReAct/Predict/ChainOfThought starts
    - on_module_end: Finalizes the run when the module completes

    All LLM calls within a module invocation are aggregated into a single run.
    Nested module calls (e.g., Predict inside ReAct) do NOT create separate runs.
    """

    def __init__(self, context: ActiveContext) -> None:
        """Initialize the callback.

        Args:
            context: The ActiveContext that holds sink, pricing, and current state
        """
        self.context = context
        self.call_id_to_run: dict[str, RunState] = {}

    def on_module_start(
        self,
        call_id: str,
        instance: Any,
        inputs: dict[str, Any],
    ) -> None:
        """Called when a DSPy module's forward() starts.

        Creates a new run only for TOP-LEVEL module calls.
        Uses ACTIVE_CALL_ID to detect if we're in a nested call.
        """
        # Only create runs for top-level module calls
        # ACTIVE_CALL_ID is None when we're at the top level
        # It gets set to the current call_id AFTER on_module_start fires
        if ACTIVE_CALL_ID.get() is not None:
            # This is a nested call (inside another module), skip
            return

        # Only track ReAct/Predict/ChainOfThought
        if not isinstance(instance, (dspy.ReAct, dspy.Predict, dspy.ChainOfThought)):
            return

        # Create new RunState for this top-level module
        run_state = RunState(
            id=generate_uuid(),
            started_at=self.context.clock.now_utc(),
            run_context_id=self.context.context_id,
            provider=self.context.provider,
            tags=list(self.context.tags),
            metadata=dict(self.context.metadata),
            input=_serialize_inputs(inputs) if inputs else None,
        )

        # POST run with status=running
        try:
            self.context.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for {type(instance).__name__}")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return

        # Store mapping and set as current
        self.call_id_to_run[call_id] = run_state
        self.context.current_run_state = run_state

        # Ensure the context is set in the current async context for @nanomon decorator
        set_active_context(self.context)

    def on_module_end(
        self,
        call_id: str,
        outputs: Any | None,
        exception: Exception | None = None,
    ) -> None:
        """Called when a DSPy module's forward() completes.

        Finalizes the run and sends the final payload.
        Only processes top-level calls that we're tracking.
        """
        if call_id not in self.call_id_to_run:
            # Not a top-level call we're tracking
            return

        run_state = self.call_id_to_run.pop(call_id)

        # Capture outputs
        run_state.output = _serialize_outputs(outputs)

        # Finalize the run
        run_state.finalize(self.context.clock.now_utc(), exception)

        # PUT with final stats
        try:
            self.context.sink.update_run(run_state.id, run_state.to_payload())
            logger.debug(
                f"Completed run {run_state.id}: status={run_state.status.value}, "
                f"tokens={run_state.input_tokens}+{run_state.output_tokens}, "
                f"cost=${run_state.cost_usd:.6f}"
            )
        except Exception as e:
            logger.error(f"Failed to update run {run_state.id}: {e}")

        # Clear current run
        if self.context.current_run_state == run_state:
            self.context.current_run_state = None
