"""Main NanomonRunContext class for nanomon."""

from __future__ import annotations

import logging
import os
from contextlib import contextmanager
from typing import Any, Callable, Iterator, TypeVar

from nanomon.decorators import create_nanomon_tracker
from nanomon.instrumentation import (
    ActiveContext,
    InstrumentedLM,
    get_active_context,
    set_active_context,
)
from nanomon.sinks import ApiSink, Sink
from nanomon.types import RunState
from nanomon.usage_extractors import ExtractorChain, UsageExtractor
from nanomon.utils import (
    Clock,
    SystemClock,
    deduplicate_tags,
    generate_uuid,
    merge_metadata,
)

F = TypeVar("F", bound=Callable[..., Any])

logger = logging.getLogger(__name__)

# Default API URL for nanomon backend (can be overridden via NANOMON_API_URL env var)
_NANOMON_DEFAULT_API_URL = "https://api-nano.nanomon.ai/api/v1"
NANOMON_API_URL = os.environ.get("NANOMON_API_URL", _NANOMON_DEFAULT_API_URL)


class NanomonRunContext:
    """Main class for tracking DSPy runs and metrics.

    Usage:
        from nanomon import NanomonRunContext

        context = NanomonRunContext(default_tags=["production"])

        lm = dspy.LM(model="gpt-4o-mini")
        dspy.settings.configure(lm=context.instrument_lm(lm))

        analyzer = dspy.ReAct(MySignature, tools=[...])

        with context.run(tags=["exp-12"], metadata={"dataset": "qa-set-1"}):
            result = await context.react(analyzer, question="Explain X")
    """

    def __init__(
        self,
        *,
        default_provider: str | None = None,
        default_tags: list[str] | None = None,
        default_metadata: dict[str, str | int | float | bool | None] | None = None,
        clock: Clock | None = None,
        sink: Sink | None = None,
    ) -> None:
        """Initialize NanomonRunContext.

        Args:
            default_provider: Default provider name for runs
            default_tags: Default tags applied to all runs
            default_metadata: Default metadata applied to all runs
            clock: Clock for timing (defaults to SystemClock)
            sink: Optional custom sink (defaults to ApiSink with nanomon backend)
        """
        self._sink = sink or ApiSink(NANOMON_API_URL)
        self._default_provider = default_provider
        self._default_tags = default_tags or []
        self._default_metadata = default_metadata or {}
        self._clock = clock or SystemClock()
        self._extractor_chain = ExtractorChain()

    def instrument_lm(self, lm: Any) -> InstrumentedLM:
        """Wrap a DSPy LM to instrument it for metrics collection.

        Args:
            lm: The DSPy LM instance to instrument

        Returns:
            An instrumented LM that behaves like the original but records metrics
        """
        return InstrumentedLM(
            lm=lm,
            extractor_chain=self._extractor_chain,
            clock=self._clock,
        )

    def register_extractor(self, extractor: UsageExtractor) -> None:
        """Register a custom usage extractor at highest priority.

        Args:
            extractor: The extractor to register
        """
        self._extractor_chain.register(extractor)

    @contextmanager
    def run(
        self,
        tags: list[str] | None = None,
        metadata: dict[str, str | int | float | bool | None] | None = None,
        provider: str | None = None,
        context_name: str | None = None,
        auto_configure_dspy: bool = True,
    ) -> Iterator[ActiveContext]:
        """Context manager for defining a run context scope.

        Automatically configures DSPy callbacks so each top-level DSPy module
        (Predict, ReAct, ChainOfThought) is tracked as a separate run.

        Args:
            tags: Tags applied to all runs in this context (merged with default_tags)
            metadata: Metadata applied to all runs (merged with default_metadata)
            provider: Provider override for runs in this context
            context_name: Optional name for the run context
            auto_configure_dspy: If True (default), automatically configure DSPy
                callbacks for run tracking. Set to False to disable.

        Yields:
            The ActiveContext object for this scope

        Example:
            with context.run(tags=["experiment-1"]):
                # DSPy callbacks are automatically configured
                predictor = dspy.Predict(MySignature)
                result = await predictor.acall(input="Hello")
                # Each DSPy module call is tracked as a separate run
        """
        merged_tags = deduplicate_tags(self._default_tags, tags)
        merged_metadata = merge_metadata(self._default_metadata, metadata)

        # Create run context
        context_id = self._sink.create_context(name=context_name, metadata=merged_metadata)
        logger.debug(f"Created run context {context_id}")

        # Create active context (no run yet - wrappers create runs)
        ctx = ActiveContext(
            context_id=context_id,
            sink=self._sink,
            clock=self._clock,
            extractor_chain=self._extractor_chain,
            tags=merged_tags,
            metadata=merged_metadata,
            provider=provider or self._default_provider,
        )

        # Save previous context (for nested contexts)
        previous_ctx = get_active_context()

        # Set this context as active
        set_active_context(ctx)

        # Auto-configure DSPy callbacks
        previous_callbacks = None
        if auto_configure_dspy:
            try:
                import dspy
                from nanomon.callback import NanomonCallback

                callback = NanomonCallback(ctx)

                # Get current callbacks (if any)
                previous_callbacks = getattr(dspy.settings, 'callbacks', None) or []

                # Add our callback (prepend so it runs first)
                dspy.configure(callbacks=[callback, *previous_callbacks])
                logger.debug("Auto-configured DSPy callback for run tracking")

            except ImportError:
                logger.debug("dspy not available, skipping auto-configure")
            except Exception as e:
                logger.warning(f"Failed to auto-configure DSPy callbacks: {e}")

        try:
            yield ctx
        finally:
            # Restore previous DSPy callbacks
            if auto_configure_dspy and previous_callbacks is not None:
                try:
                    import dspy
                    dspy.configure(callbacks=previous_callbacks)
                    logger.debug("Restored previous DSPy callbacks")
                except Exception:
                    pass

            set_active_context(previous_ctx)
            logger.debug(f"Closed run context {context_id}")

    def _create_run_state(self, ctx: ActiveContext) -> RunState:
        """Create a new RunState for tracking a module execution."""
        return RunState(
            id=generate_uuid(),
            started_at=self._clock.now_utc(),
            run_context_id=ctx.context_id,
            provider=ctx.provider,
            tags=list(ctx.tags),
            metadata=dict(ctx.metadata),
        )

    async def react(self, module: Any, **inputs: Any) -> Any:
        """Execute a dspy.ReAct module as a tracked run.

        Creates a run that aggregates all internal LLM calls and tracks
        tool calls made during execution.

        Args:
            module: The dspy.ReAct instance to execute
            **inputs: Arguments to pass to module.aforward()

        Returns:
            The module's output (dspy.Prediction)

        Example:
            analyzer = dspy.ReAct(CodeAnalysisSignature, tools=fs_tools)
            result = await context.react(
                analyzer,
                task_title="Implement feature X",
                task_description="...",
            )
        """
        ctx = get_active_context()
        if ctx is None:
            # No context, just call without tracking
            return await module.aforward(**inputs)

        run_state = self._create_run_state(ctx)

        # POST run with status=running
        try:
            ctx.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for ReAct")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return await module.aforward(**inputs)

        # Set current run so tools and LLM calls can attach
        ctx.current_run_state = run_state

        error = None
        try:
            return await module.aforward(**inputs)
        except Exception as e:
            error = e
            raise
        finally:
            run_state.finalize(self._clock.now_utc(), error)
            try:
                ctx.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}"
                )
            except Exception as e:
                logger.error(f"Failed to update run {run_state.id}: {e}")
            ctx.current_run_state = None

    async def predict(self, module: Any, **inputs: Any) -> Any:
        """Execute a dspy.Predict module as a tracked run.

        Creates a run that tracks the LLM call and any tool calls.

        Args:
            module: The dspy.Predict instance to execute
            **inputs: Arguments to pass to module.aforward()

        Returns:
            The module's output (dspy.Prediction)

        Example:
            predictor = dspy.Predict(SummarizationSignature)
            result = await context.predict(predictor, text="...")
        """
        ctx = get_active_context()
        if ctx is None:
            return await module.aforward(**inputs)

        run_state = self._create_run_state(ctx)

        try:
            ctx.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for Predict")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return await module.aforward(**inputs)

        ctx.current_run_state = run_state

        error = None
        try:
            return await module.aforward(**inputs)
        except Exception as e:
            error = e
            raise
        finally:
            run_state.finalize(self._clock.now_utc(), error)
            try:
                ctx.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}"
                )
            except Exception as e:
                logger.error(f"Failed to update run {run_state.id}: {e}")
            ctx.current_run_state = None

    async def module(self, module: Any, **inputs: Any) -> Any:
        """Execute any DSPy module as a tracked run.

        Generic wrapper that works with any DSPy module (ReAct, Predict,
        ChainOfThought, custom modules, etc.).

        Args:
            module: The DSPy module instance to execute
            **inputs: Arguments to pass to module.aforward()

        Returns:
            The module's output

        Example:
            custom_module = MyCustomDSPyModule()
            result = await context.module(custom_module, input="...")
        """
        ctx = get_active_context()
        if ctx is None:
            return await module.aforward(**inputs)

        run_state = self._create_run_state(ctx)

        try:
            ctx.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for {type(module).__name__}")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return await module.aforward(**inputs)

        ctx.current_run_state = run_state

        error = None
        try:
            return await module.aforward(**inputs)
        except Exception as e:
            error = e
            raise
        finally:
            run_state.finalize(self._clock.now_utc(), error)
            try:
                ctx.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}"
                )
            except Exception as e:
                logger.error(f"Failed to update run {run_state.id}: {e}")
            ctx.current_run_state = None

    def react_sync(self, module: Any, **inputs: Any) -> Any:
        """Synchronous version of react() for non-async contexts."""
        ctx = get_active_context()
        if ctx is None:
            return module(**inputs)

        run_state = self._create_run_state(ctx)

        try:
            ctx.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for ReAct (sync)")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return module(**inputs)

        ctx.current_run_state = run_state

        error = None
        try:
            return module(**inputs)
        except Exception as e:
            error = e
            raise
        finally:
            run_state.finalize(self._clock.now_utc(), error)
            try:
                ctx.sink.update_run(run_state.id, run_state.to_payload())
            except Exception as e:
                logger.error(f"Failed to update run {run_state.id}: {e}")
            ctx.current_run_state = None

    def predict_sync(self, module: Any, **inputs: Any) -> Any:
        """Synchronous version of predict() for non-async contexts."""
        ctx = get_active_context()
        if ctx is None:
            return module(**inputs)

        run_state = self._create_run_state(ctx)

        try:
            ctx.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for Predict (sync)")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return module(**inputs)

        ctx.current_run_state = run_state

        error = None
        try:
            return module(**inputs)
        except Exception as e:
            error = e
            raise
        finally:
            run_state.finalize(self._clock.now_utc(), error)
            try:
                ctx.sink.update_run(run_state.id, run_state.to_payload())
            except Exception as e:
                logger.error(f"Failed to update run {run_state.id}: {e}")
            ctx.current_run_state = None

    def module_sync(self, module: Any, **inputs: Any) -> Any:
        """Synchronous version of module() for non-async contexts."""
        ctx = get_active_context()
        if ctx is None:
            return module(**inputs)

        run_state = self._create_run_state(ctx)

        try:
            ctx.sink.write_run(run_state.to_payload())
            logger.debug(f"Created run {run_state.id} for {type(module).__name__} (sync)")
        except Exception as e:
            logger.error(f"Failed to create run: {e}")
            return module(**inputs)

        ctx.current_run_state = run_state

        error = None
        try:
            return module(**inputs)
        except Exception as e:
            error = e
            raise
        finally:
            run_state.finalize(self._clock.now_utc(), error)
            try:
                ctx.sink.update_run(run_state.id, run_state.to_payload())
            except Exception as e:
                logger.error(f"Failed to update run {run_state.id}: {e}")
            ctx.current_run_state = None

    def nanomon(
        self,
        name: str | None = None,
        method: str | None = None,
    ) -> Callable[[F], F]:
        """Decorator to track tool calls within a run context.

        Use this decorator on functions that represent tool calls in your
        agent/ReAct workflow. Tool calls will be recorded with timing and
        success/error status.

        Args:
            name: Optional name override for the tool (defaults to function name)
            method: Optional method name for additional categorization

        Returns:
            A decorator that tracks the tool call

        Example:
            context = NanomonRunContext()

            @context.nanomon(name="web_search")
            def search_web(query: str) -> str:
                return results

            with context.run():
                await context.react(analyzer, query="...")
                # Tool calls during ReAct execution are automatically tracked
        """
        return create_nanomon_tracker(self._sink, self._clock)(name, method)
