"""Main RunLogger class for nanomon."""

from __future__ import annotations

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

from nanomon.decorators import create_tool_tracker
from nanomon.instrumentation import (
    ActiveContext,
    InstrumentedLM,
    get_active_context,
    set_active_context,
)
from nanomon.pricing import PricingSource
from nanomon.sinks import 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])

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


class RunLogger:
    """Main class for logging DSPy run costs and metrics.

    Usage:
        from nanomon import RunLogger
        from nanomon.sinks import ApiSink
        from nanomon.pricing import JsonPricing

        pricing = JsonPricing("pricing.json")
        sink = ApiSink("http://localhost:8000/api/v1")
        cost_logger = RunLogger(pricing=pricing, sink=sink, default_tags=["local"])

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

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

        with cost_logger.run(tags=["exp-12"], metadata={"dataset": "qa-set-1"}):
            # Use explicit wrapper to create a tracked run
            result = await cost_logger.react(analyzer, question="Explain X")
    """

    def __init__(
        self,
        pricing: PricingSource,
        sink: Sink,
        *,
        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,
    ) -> None:
        """Initialize the RunLogger.

        Args:
            pricing: Pricing source for cost calculation
            sink: Sink for persisting runs
            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)
        """
        self._pricing = pricing
        self._sink = sink
        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,
            pricing=self._pricing,
            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,
    ) -> Iterator[ActiveContext]:
        """Context manager for defining a run context scope.

        Use react(), predict(), or module() wrapper methods within this context
        to create tracked runs. Each wrapper call creates one run with all its
        internal LLM calls aggregated.

        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

        Yields:
            The ActiveContext object for this scope

        Example:
            with cost_logger.run(tags=["experiment-1"]) as ctx:
                result = await cost_logger.react(analyzer, question="What is 2+2?")
                # Creates one run with all internal LLM calls aggregated
        """
        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,
            pricing=self._pricing,
            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)

        try:
            yield ctx
        finally:
            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 cost_logger.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}, "
                    f"cost=${run_state.cost_usd:.6f}"
                )
            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 cost_logger.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}, "
                    f"cost=${run_state.cost_usd:.6f}"
                )
            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 cost_logger.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}, "
                    f"cost=${run_state.cost_usd:.6f}"
                )
            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 track_tool(
        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:
            cost_logger = RunLogger(pricing=..., sink=...)

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

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