"""
LangChain auto-injection patches.

Patches langchain_core.runnables.base.Runnable to automatically inject
RisicareCallbackHandler into chain executions. Also sets suppression
context to prevent double-tracing with provider patches.

Patched methods:
    Runnable.invoke    → sync execution
    Runnable.ainvoke   → async execution
    Runnable.batch     → batch sync execution
    Runnable.abatch    → batch async execution
"""

from __future__ import annotations

import logging
from typing import Any, Callable, Dict, TypeVar

import wrapt

from risicare.integrations._base import is_tracing_enabled
from risicare.integrations._dedup import (
    suppress_provider_instrumentation,
    async_suppress_provider_instrumentation,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")

_singleton_handler: Any = None


def _get_callback_handler() -> Any:
    """Get or create the singleton RisicareCallbackHandler instance.

    Using a singleton ensures all chain invocations share the same
    _run_spans dict, so parent-child linking works across nested chains
    and orphan cleanup applies globally.
    """
    global _singleton_handler
    if _singleton_handler is None:
        from risicare.integrations.langchain._callback import RisicareCallbackHandler

        _singleton_handler = RisicareCallbackHandler()
    return _singleton_handler


def _inject_callback(kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Inject RisicareCallbackHandler into the config's callbacks list.

    Handles the LangChain config pattern where callbacks can live in
    kwargs["config"]["callbacks"] or kwargs["callbacks"].

    Returns the modified kwargs dict.
    """
    handler = _get_callback_handler()

    # LangChain uses a RunnableConfig dict for configuration
    config = kwargs.get("config")
    if config is None:
        config = {}
        kwargs["config"] = config

    callbacks = config.get("callbacks")
    if callbacks is None:
        callbacks = []
        config["callbacks"] = callbacks

    # Avoid duplicate injection (check by type)
    handler_type = type(handler)
    if not any(isinstance(cb, handler_type) for cb in callbacks):
        callbacks.append(handler)

    return kwargs


def _wrap_invoke(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Runnable.invoke that injects callback + suppresses providers."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    kwargs = _inject_callback(kwargs)
    with suppress_provider_instrumentation():
        return wrapped(*args, **kwargs)


async def _wrap_ainvoke(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Runnable.ainvoke that injects callback + suppresses providers."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    kwargs = _inject_callback(kwargs)
    async with async_suppress_provider_instrumentation():
        return await wrapped(*args, **kwargs)


def _wrap_batch(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Runnable.batch that injects callback + suppresses providers."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    kwargs = _inject_callback(kwargs)
    with suppress_provider_instrumentation():
        return wrapped(*args, **kwargs)


async def _wrap_abatch(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Runnable.abatch that injects callback + suppresses providers."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    kwargs = _inject_callback(kwargs)
    async with async_suppress_provider_instrumentation():
        return await wrapped(*args, **kwargs)


def patch_langchain(module: Any) -> None:
    """
    Apply wrapt patches to LangChain Runnable methods.

    Patches are applied to langchain_core.runnables.base.Runnable to
    capture all chain execution regardless of which specific Runnable
    subclass is used.
    """
    try:
        wrapt.wrap_function_wrapper(
            "langchain_core.runnables.base",
            "Runnable.invoke",
            _wrap_invoke,
        )
    except Exception as e:
        logger.debug(f"Failed to patch Runnable.invoke: {e}")

    try:
        wrapt.wrap_function_wrapper(
            "langchain_core.runnables.base",
            "Runnable.ainvoke",
            _wrap_ainvoke,
        )
    except Exception as e:
        logger.debug(f"Failed to patch Runnable.ainvoke: {e}")

    try:
        wrapt.wrap_function_wrapper(
            "langchain_core.runnables.base",
            "Runnable.batch",
            _wrap_batch,
        )
    except Exception as e:
        logger.debug(f"Failed to patch Runnable.batch: {e}")

    try:
        wrapt.wrap_function_wrapper(
            "langchain_core.runnables.base",
            "Runnable.abatch",
            _wrap_abatch,
        )
    except Exception as e:
        logger.debug(f"Failed to patch Runnable.abatch: {e}")
