"""LangChain auto-instrumentation.

Patches ``langchain_core.language_models.chat_models.BaseChatModel.invoke``
and ``ainvoke`` as well as chain/agent ``invoke``/``ainvoke`` to capture
model_call spans with token counts and cost.
"""
from __future__ import annotations

import functools
import logging
from datetime import datetime, timezone
from typing import Any

from fibinaci.instrumentation._base import BaseInstrumentor
from fibinaci.instrumentation._wrapper import ProviderConfig, _make_error_span

logger = logging.getLogger("fibinaci.instrumentation")

_CONFIG = ProviderConfig(
    name="langchain",
    input_tokens_field="prompt_tokens",
    output_tokens_field="completion_tokens",
    finish_reason_getter=lambda r: getattr(r, "response_metadata", {}).get("finish_reason"),
)


def _extract_langchain_usage(response: Any) -> tuple[int | None, int | None, str | None]:
    """Pull token counts and model from a LangChain AIMessage or similar."""
    usage = getattr(response, "usage_metadata", None)
    if usage:
        inp = usage.get("input_tokens") if isinstance(usage, dict) else getattr(usage, "input_tokens", None)
        out = usage.get("output_tokens") if isinstance(usage, dict) else getattr(usage, "output_tokens", None)
    else:
        meta = getattr(response, "response_metadata", {}) or {}
        token_usage = meta.get("token_usage", {})
        inp = token_usage.get("prompt_tokens")
        out = token_usage.get("completion_tokens")

    model = None
    meta = getattr(response, "response_metadata", {}) or {}
    model = meta.get("model_name") or meta.get("model")
    return inp, out, model


def _wrap_invoke(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData, estimate_cost

                inp, out, model = _extract_langchain_usage(response)
                model = model or getattr(self, "model_name", None) or getattr(self, "model", None)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="langchain",
                    model_type="llm",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                    cost_usd=estimate_cost(model, inp, out),
                )
                sd.set_model_call(model_name=model, finish_reason=_CONFIG.finish_reason_getter(response))
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture langchain span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                model = getattr(self, "model_name", None) or getattr(self, "model", None)
                get_buffer().push(_make_error_span(_CONFIG, {"model": model}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture langchain error span", exc_info=True)
            raise

    return wrapper


def _wrap_ainvoke(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData, estimate_cost

                inp, out, model = _extract_langchain_usage(response)
                model = model or getattr(self, "model_name", None) or getattr(self, "model", None)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="langchain",
                    model_type="llm",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                    cost_usd=estimate_cost(model, inp, out),
                )
                sd.set_model_call(model_name=model, finish_reason=_CONFIG.finish_reason_getter(response))
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async langchain span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                model = getattr(self, "model_name", None) or getattr(self, "model", None)
                get_buffer().push(_make_error_span(_CONFIG, {"model": model}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async langchain error span", exc_info=True)
            raise

    return wrapper


class LangChainInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._originals: dict[str, Any] = {}

    def instrument(self) -> None:
        from langchain_core.language_models.chat_models import BaseChatModel

        self._originals["chat_invoke"] = BaseChatModel.invoke
        BaseChatModel.invoke = _wrap_invoke(BaseChatModel.invoke)

        self._originals["chat_ainvoke"] = BaseChatModel.ainvoke
        BaseChatModel.ainvoke = _wrap_ainvoke(BaseChatModel.ainvoke)

    def uninstrument(self) -> None:
        if not self._originals:
            return
        from langchain_core.language_models.chat_models import BaseChatModel

        BaseChatModel.invoke = self._originals["chat_invoke"]
        BaseChatModel.ainvoke = self._originals["chat_ainvoke"]
        self._originals.clear()
