# Copyright The OpenTelemetry Authors
# SPDX-License-Identifier: Apache-2.0

"""
Langchain instrumentation supporting `ChatOpenAI` and `ChatBedrock`, it can be enabled by
using ``LangChainInstrumentor``. Other providers/LLMs may be supported in the future and telemetry for them is skipped for now.

Usage
-----
.. code:: python
    from opentelemetry.instrumentation.langchain import LangChainInstrumentor
    from langchain_core.messages import HumanMessage, SystemMessage
    from langchain_openai import ChatOpenAI

    LangChainInstrumentor().instrument()
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, max_tokens=1000)
    messages = [
        SystemMessage(content="You are a helpful assistant!"),
        HumanMessage(content="What is the capital of France?"),
    ]
    result = llm.invoke(messages)
    LangChainInstrumentor().uninstrument()

API
---
"""

from typing import Any, Callable, Collection

from langchain_core.callbacks import BaseCallbackHandler
from wrapt import wrap_function_wrapper  # type: ignore

from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.langchain.callback_handler import (
    OpenTelemetryLangChainCallbackHandler,
)
from opentelemetry.instrumentation.langchain.package import _instruments
from opentelemetry.instrumentation.utils import unwrap
from opentelemetry.util.genai.handler import get_telemetry_handler


class LangChainInstrumentor(BaseInstrumentor):
    """
    OpenTelemetry instrumentor for LangChain.
    This adds a custom callback handler to the LangChain callback manager
    to capture LLM telemetry.
    """

    def __init__(
        self,
    ):
        super().__init__()

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs: Any):
        """
        Enable Langchain instrumentation.
        """
        tracer_provider = kwargs.get("tracer_provider")
        meter_provider = kwargs.get("meter_provider")
        logger_provider = kwargs.get("logger_provider")

        telemetry_handler = get_telemetry_handler(
            tracer_provider=tracer_provider,
            meter_provider=meter_provider,
            logger_provider=logger_provider,
        )
        otel_callback_handler = OpenTelemetryLangChainCallbackHandler(
            telemetry_handler=telemetry_handler,
        )

        wrap_function_wrapper(
            "langchain_core.callbacks",
            "BaseCallbackManager.__init__",
            _BaseCallbackManagerInitWrapper(otel_callback_handler),
        )

    def _uninstrument(self, **kwargs: Any):
        """
        Cleanup instrumentation (unwrap).
        """
        unwrap("langchain_core.callbacks.base.BaseCallbackManager", "__init__")
        # Clear the TelemetryHandler singleton so the next instrument() uses
        # the provided tracer_provider/meter_provider/logger_provider instead
        # of reusing the previous handler.
        if (
            getattr(get_telemetry_handler, "_default_handler", None)
            is not None
        ):
            delattr(get_telemetry_handler, "_default_handler")


class _BaseCallbackManagerInitWrapper:
    """
    Wrap the BaseCallbackManager __init__ to insert custom callback handler in the manager's handlers list.
    """

    def __init__(
        self, callback_handler: OpenTelemetryLangChainCallbackHandler
    ):
        self._otel_handler = callback_handler

    def __call__(
        self,
        wrapped: Callable[..., None],
        instance: BaseCallbackHandler,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ):
        wrapped(*args, **kwargs)
        # Ensure our OTel callback is present if not already.
        for handler in instance.inheritable_handlers:  # type: ignore
            if isinstance(handler, type(self._otel_handler)):
                break
        else:
            instance.add_handler(self._otel_handler, inherit=True)  # type: ignore
