"""Decorator factories for nanomon."""

from __future__ import annotations

import asyncio
import inspect
import logging
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, TypeVar

from nanomon.instrumentation import get_active_context
from nanomon.utils import Clock, SystemClock

if TYPE_CHECKING:
    from nanomon.sinks import Sink

logger = logging.getLogger(__name__)

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


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()}
    return str(value)  # Fallback to string representation


def _capture_tool_input(func: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any] | None:
    """Capture function arguments as a dictionary for tool input tracking."""
    input_data: dict[str, Any] = {}

    # Get parameter names from function signature
    try:
        sig = inspect.signature(func)
        params = list(sig.parameters.keys())

        # Map positional args to parameter names
        for i, arg in enumerate(args):
            if i < len(params):
                input_data[params[i]] = _serialize_value(arg)
            else:
                input_data[f"arg_{i}"] = _serialize_value(arg)

        # Add keyword arguments
        for key, value in kwargs.items():
            input_data[key] = _serialize_value(value)

    except (ValueError, TypeError):
        # Fallback if signature inspection fails
        for i, arg in enumerate(args):
            input_data[f"arg_{i}"] = _serialize_value(arg)
        for key, value in kwargs.items():
            input_data[key] = _serialize_value(value)

    return input_data if input_data else None


def _get_tool_name(func: Callable[..., Any], explicit_name: str | None) -> str:
    """Get tool name, using explicit name or deriving from function.

    If no explicit name provided, returns '<module_basename>.<func_name>'.
    """
    if explicit_name:
        return explicit_name
    module = func.__module__ or "unknown"
    # Get just the last part of the module path (e.g., "filesystem_tools" from "tools.filesystem_tools")
    module_basename = module.rsplit(".", 1)[-1]
    return f"{module_basename}.{func.__name__}"


# Module-level sink/clock for standalone @nanomon decorator
_default_sink: Sink | None = None
_default_clock: Clock = SystemClock()


def configure_nanomon(sink: Sink, clock: Clock | None = None) -> None:
    """Configure the default sink and clock for standalone @nanomon decorator.

    Call this after creating your sink to enable tool tracking with the
    standalone @nanomon decorator.

    Args:
        sink: The sink to write tool calls to
        clock: Optional clock for timing (defaults to SystemClock)

    Example:
        from nanomon import ApiSink
        from nanomon.decorators import configure_nanomon

        sink = ApiSink("http://localhost:8000/api/v1")
        configure_nanomon(sink)

        # Now @nanomon decorators will work
    """
    global _default_sink, _default_clock
    _default_sink = sink
    _default_clock = clock or SystemClock()


def create_nanomon_tracker(sink: Sink, clock: Clock) -> Callable[..., Callable[[F], F]]:
    """Factory that creates the @nanomon decorator with sink/clock access.

    Args:
        sink: The sink to write tool calls to
        clock: The clock for timing

    Returns:
        A nanomon decorator function
    """

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

        Args:
            name: Optional name override for the tool (defaults to <module>.<func_name>)
            method: Optional method name for additional categorization

        Returns:
            A decorator that tracks the tool call
        """

        def decorator(func: F) -> F:
            tool_name = _get_tool_name(func, name)

            @wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                ctx = get_active_context()
                if ctx is None or ctx.current_run_id is None:
                    # No active run, just call the function
                    return func(*args, **kwargs)

                run_id = ctx.current_run_id
                started_at = clock.now_utc()
                status = "success"
                error_msg: str | None = None
                tool_input = _capture_tool_input(func, args, kwargs)

                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    status = "error"
                    error_msg = str(e)
                    raise
                finally:
                    ended_at = clock.now_utc()
                    payload = {
                        "name": tool_name,
                        "method": method,
                        "startedAt": started_at.isoformat().replace("+00:00", "Z"),
                        "endedAt": ended_at.isoformat().replace("+00:00", "Z"),
                        "status": status,
                        "error": error_msg,
                        "input": tool_input,
                    }
                    try:
                        sink.write_tool_call(run_id, payload)
                        logger.debug(f"Tracked tool call: {tool_name}")
                    except Exception as e:
                        logger.warning(f"Failed to write tool call {tool_name}: {e}")

            @wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                ctx = get_active_context()
                if ctx is None or ctx.current_run_id is None:
                    # No active run, just call the function
                    return await func(*args, **kwargs)

                run_id = ctx.current_run_id
                started_at = clock.now_utc()
                status = "success"
                error_msg: str | None = None
                tool_input = _capture_tool_input(func, args, kwargs)

                try:
                    return await func(*args, **kwargs)
                except Exception as e:
                    status = "error"
                    error_msg = str(e)
                    raise
                finally:
                    ended_at = clock.now_utc()
                    payload = {
                        "name": tool_name,
                        "method": method,
                        "startedAt": started_at.isoformat().replace("+00:00", "Z"),
                        "endedAt": ended_at.isoformat().replace("+00:00", "Z"),
                        "status": status,
                        "error": error_msg,
                        "input": tool_input,
                    }
                    try:
                        sink.write_tool_call(run_id, payload)
                        logger.debug(f"Tracked tool call: {tool_name}")
                    except Exception as e:
                        logger.warning(f"Failed to write tool call {tool_name}: {e}")

            # Return the appropriate wrapper based on function type
            if asyncio.iscoroutinefunction(func):
                return async_wrapper  # type: ignore[return-value]
            return sync_wrapper  # type: ignore[return-value]

        return decorator

    return nanomon


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

    This decorator can be used at module definition time, before a NanomonRunContext
    is created. Call configure_nanomon() with your sink before running
    the decorated functions.

    If no sink is configured or no run is active, the function executes
    normally without tracking.

    Args:
        name: Optional name override for the tool (defaults to <module>.<func_name>)
        method: Optional method name for additional categorization

    Returns:
        A decorator that tracks the tool call

    Example:
        from nanomon.decorators import nanomon, configure_nanomon

        @nanomon()  # Will be tracked as "my_module.search_web"
        def search_web(query: str) -> str:
            return results

        @nanomon(name="custom_name")  # Explicit name override
        def another_tool() -> str:
            return results

        # Later, when setting up the logger:
        configure_nanomon(sink)

        with context.run():
            search_web("query")  # Now tracked
    """

    def decorator(func: F) -> F:
        tool_name = _get_tool_name(func, name)

        @wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            ctx = get_active_context()
            if ctx is None or ctx.current_run_id is None or _default_sink is None:
                # No active run or no sink configured, just call the function
                return func(*args, **kwargs)

            run_id = ctx.current_run_id
            started_at = _default_clock.now_utc()
            status = "success"
            error_msg: str | None = None
            tool_input = _capture_tool_input(func, args, kwargs)

            try:
                return func(*args, **kwargs)
            except Exception as e:
                status = "error"
                error_msg = str(e)
                raise
            finally:
                ended_at = _default_clock.now_utc()
                payload = {
                    "name": tool_name,
                    "method": method,
                    "startedAt": started_at.isoformat().replace("+00:00", "Z"),
                    "endedAt": ended_at.isoformat().replace("+00:00", "Z"),
                    "status": status,
                    "error": error_msg,
                    "input": tool_input,
                }
                try:
                    _default_sink.write_tool_call(run_id, payload)
                    logger.debug(f"Tracked tool call: {tool_name}")
                except Exception as e:
                    logger.warning(f"Failed to write tool call {tool_name}: {e}")

        @wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            ctx = get_active_context()
            if ctx is None or ctx.current_run_id is None or _default_sink is None:
                # No active run or no sink configured, just call the function
                return await func(*args, **kwargs)

            run_id = ctx.current_run_id
            started_at = _default_clock.now_utc()
            status = "success"
            error_msg: str | None = None
            tool_input = _capture_tool_input(func, args, kwargs)

            try:
                return await func(*args, **kwargs)
            except Exception as e:
                status = "error"
                error_msg = str(e)
                raise
            finally:
                ended_at = _default_clock.now_utc()
                payload = {
                    "name": tool_name,
                    "method": method,
                    "startedAt": started_at.isoformat().replace("+00:00", "Z"),
                    "endedAt": ended_at.isoformat().replace("+00:00", "Z"),
                    "status": status,
                    "error": error_msg,
                    "input": tool_input,
                }
                try:
                    _default_sink.write_tool_call(run_id, payload)
                    logger.debug(f"Tracked tool call: {tool_name}")
                except Exception as e:
                    logger.warning(f"Failed to write tool call {tool_name}: {e}")

        # Return the appropriate wrapper based on function type
        if asyncio.iscoroutinefunction(func):
            return async_wrapper  # type: ignore[return-value]
        return sync_wrapper  # type: ignore[return-value]

    return decorator
