"""Core wrapping logic shared by ``instrument()`` and ``@langfuse_trace``.

This module is internal.  Public API is re-exported from ``__init__.py``.
"""

from __future__ import annotations

import asyncio
import functools
import inspect
import json
from typing import Any, Callable, TypeVar

from opentelemetry import trace as otel_trace

from ._enrichment import _enrich_current_span, _tag_span_io
from .config import LangfuseConfig

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

_MAX_IO_LEN = 2000


def _safe_json(obj: Any) -> str:
    """Best-effort JSON serialisation, truncated to *_MAX_IO_LEN*."""
    try:
        raw = json.dumps(obj, default=str)
    except (TypeError, ValueError):
        raw = str(obj)
    return raw if len(raw) <= _MAX_IO_LEN else raw[:_MAX_IO_LEN] + "…"


def _serialize_args(fn: Callable, args: tuple, kwargs: dict) -> str:
    """Bind ``*args`` / ``**kwargs`` to *fn*'s signature -> JSON string."""
    try:
        sig = inspect.signature(fn)
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()
        return _safe_json(dict(bound.arguments))
    except Exception:
        return _safe_json({"args": list(args), "kwargs": kwargs})


def _wrap_handler(fn: F, config: LangfuseConfig) -> F:
    """Wrap a handler with automatic Langfuse enrichment and I/O capture.

    Works for both sync and async handlers.
    """
    if asyncio.iscoroutinefunction(fn):

        @functools.wraps(fn)
        async def _async(*args: Any, **kwargs: Any) -> Any:
            span = otel_trace.get_current_span()
            _enrich_current_span(config)
            input_json = _serialize_args(fn, args, kwargs)
            _tag_span_io(span, input=input_json)
            result = await fn(*args, **kwargs)
            _tag_span_io(span, input=input_json, output=_safe_json(result))
            return result

        return _async

    @functools.wraps(fn)
    def _sync(*args: Any, **kwargs: Any) -> Any:
        span = otel_trace.get_current_span()
        _enrich_current_span(config)
        input_json = _serialize_args(fn, args, kwargs)
        _tag_span_io(span, input=input_json)
        result = fn(*args, **kwargs)
        _tag_span_io(span, input=input_json, output=_safe_json(result))
        return result

    return _sync


def _patch_registration(mcp: Any, method_name: str, config: LangfuseConfig) -> None:
    """Monkey-patch one registration method so it auto-wraps handlers."""
    original = getattr(mcp, method_name)

    @functools.wraps(original)
    def patched(*args: Any, **kwargs: Any) -> Callable:
        decorator = original(*args, **kwargs)

        @functools.wraps(decorator)
        def wrapping_decorator(fn: F) -> F:
            return decorator(_wrap_handler(fn, config))

        return wrapping_decorator

    setattr(mcp, method_name, patched)
