"""LangChain BaseCallbackHandler implementation for Mingx tracing."""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence, Union
from uuid import UUID

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.documents import Document
    from langchain_core.messages import (
        AIMessage,
        BaseMessage,
        ChatMessage,
        FunctionMessage,
        HumanMessage,
        SystemMessage,
        ToolMessage,
    )
    from langchain_core.outputs import ChatGeneration, LLMResult
except ImportError as e:
    raise ModuleNotFoundError(
        "Mingx LangChain integration requires langchain-core. "
        "Install with: pip install mingx[langchain] or pip install langchain-core>=0.3.0"
    ) from e

from opentelemetry import context, trace
from opentelemetry.context import _RUNTIME_CONTEXT
from mingx._client.attributes import MingxOtelSpanAttributes
from mingx._client.constants import ObservationType, normalize_observation_type
from mingx._client.get_client import get_client
from mingx._client.span import MingxEvent, MingxGeneration, MingxSpan


def _get_run_name(serialized: Optional[Dict[str, Any]], **kwargs: Any) -> str:
    if kwargs.get("name") is not None:
        return str(kwargs["name"])
    if serialized is None:
        return "run"
    if isinstance(serialized.get("name"), str):
        return serialized["name"]
    ids = serialized.get("id")
    if isinstance(ids, list) and ids:
        return str(ids[-1])
    return "run"


def _get_observation_type_from_serialized(
    serialized: Optional[Dict[str, Any]],
    callback_type: str,
    **kwargs: Any,
) -> Union[ObservationType, str]:
    """根据 LangChain 回调类型与 serialized 推断观测类型（与 Langfuse 规则一致）。"""
    if callback_type == "tool":
        return ObservationType.TOOL
    if callback_type == "retriever":
        return ObservationType.RETRIEVER
    if callback_type == "llm":
        return ObservationType.GENERATION
    if callback_type == "chain":
        if serialized and "id" in serialized:
            class_path = serialized["id"]
            if isinstance(class_path, list):
                if any("agent" in str(part).lower() for part in class_path):
                    return ObservationType.AGENT
            elif "agent" in str(class_path).lower():
                return ObservationType.AGENT
        name = _get_run_name(serialized, **kwargs)
        if "agent" in name.lower():
            return ObservationType.AGENT
        return ObservationType.CHAIN
    return ObservationType.SPAN


def _convert_message_to_dict(msg: BaseMessage) -> Dict[str, Any]:
    """将 LangChain BaseMessage 转为统一 dict（与 Langfuse 格式一致，便于 OTLP 后端解析）。"""
    content = getattr(msg, "content", "")
    if isinstance(msg, HumanMessage):
        out: Dict[str, Any] = {"role": "user", "content": content}
    elif isinstance(msg, AIMessage):
        out = {"role": "assistant", "content": content}
        if getattr(msg, "tool_calls", None):
            out["tool_calls"] = msg.tool_calls
    elif isinstance(msg, SystemMessage):
        out = {"role": "system", "content": content}
    elif isinstance(msg, ToolMessage):
        out = {
            "role": "tool",
            "content": content,
            "tool_call_id": getattr(msg, "tool_call_id", None),
        }
    elif isinstance(msg, FunctionMessage):
        out = {"role": "function", "content": content}
    elif isinstance(msg, ChatMessage):
        out = {"role": getattr(msg, "role", "message"), "content": content}
    else:
        out = {"role": getattr(msg, "type", "message"), "content": content}
    if getattr(msg, "additional_kwargs", None) and "name" in msg.additional_kwargs:
        out["name"] = msg.additional_kwargs["name"]
    return out


def _normalize_usage_details(raw: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    """将各厂商的 usage 字段统一为 prompt_tokens / completion_tokens / total_tokens（与 Langfuse 一致）。"""
    if not raw or not isinstance(raw, dict):
        return raw
    result = dict(raw)
    mapping = [
        ("input_tokens", "prompt_tokens"),
        ("output_tokens", "completion_tokens"),
        ("input_token_count", "prompt_tokens"),
        ("generated_token_count", "completion_tokens"),
        ("prompt_token_count", "prompt_tokens"),
        ("candidates_token_count", "completion_tokens"),
        ("inputTokenCount", "prompt_tokens"),
        ("outputTokenCount", "completion_tokens"),
        ("totalTokenCount", "total_tokens"),
    ]
    for old_key, new_key in mapping:
        if old_key in result and result[old_key] is not None and new_key not in result:
            val = result[old_key]
            result[new_key] = sum(val) if isinstance(val, list) else val
    out = {k: result[k] for k in ("prompt_tokens", "completion_tokens", "total_tokens") if result.get(k) is not None}
    return out if out else (result if result else None)


def _parse_usage_from_response(response: LLMResult) -> Optional[Dict[str, Any]]:
    """从 LLMResult 中收集 usage（与 Langfuse 一致：llm_output.token_usage/usage、generation_info.usage_metadata、message.usage_metadata 等）。"""
    usage = None
    if response.llm_output:
        for key in ("token_usage", "usage"):
            if key in response.llm_output and response.llm_output[key]:
                usage = _normalize_usage_details(
                    response.llm_output[key] if isinstance(response.llm_output[key], dict) else None
                )
                if usage:
                    return usage
    if getattr(response, "generations", None):
        for gen_list in response.generations:
            for g in gen_list or []:
                info = getattr(g, "generation_info", None) or {}
                if isinstance(info, dict) and info.get("usage_metadata"):
                    usage = _normalize_usage_details(info["usage_metadata"])
                    if usage:
                        return usage
                msg = getattr(g, "message", None)
                if msg is not None:
                    um = getattr(msg, "usage_metadata", None)
                    if um:
                        usage = _normalize_usage_details(um if isinstance(um, dict) else None)
                        if usage:
                            return usage
                    rm = getattr(msg, "response_metadata", None)
                    if isinstance(rm, dict):
                        for ukey in ("usage", "amazon-bedrock-invocationMetrics"):
                            if rm.get(ukey):
                                usage = _normalize_usage_details(rm[ukey] if isinstance(rm[ukey], dict) else None)
                                if usage:
                                    return usage
    return usage


class MingxCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler that sends traces to Mingx (OTLP)."""

    def __init__(self, *, client: Optional[Any] = None) -> None:
        self._client = client or get_client()
        self._runs: Dict[UUID, Union[MingxSpan, MingxGeneration]] = {}
        self._context_tokens: Dict[UUID, Any] = {}  # token from context.attach

    def _get_parent(self, parent_run_id: Optional[UUID]) -> Union[Any, MingxSpan, MingxGeneration]:
        if parent_run_id is not None and parent_run_id in self._runs:
            return self._runs[parent_run_id]
        return self._client

    def _attach(self, run_id: UUID, observation: Union[MingxSpan, MingxGeneration]) -> None:
        ctx = trace.set_span_in_context(observation._otel_span)
        token = context.attach(ctx)
        self._runs[run_id] = observation
        self._context_tokens[run_id] = token

    def _detach(self, run_id: UUID) -> Optional[Union[MingxSpan, MingxGeneration]]:
        obs = self._runs.pop(run_id, None)
        token = self._context_tokens.pop(run_id, None)
        if token is not None:
            try:
                _RUNTIME_CONTEXT.detach(token)
            except Exception:
                pass
        return obs

    def _start_obs(
        self,
        parent_run_id: Optional[UUID],
        name: str,
        as_type: Union[ObservationType, str] = ObservationType.SPAN,
        input: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **extra: Any,
    ) -> Optional[Union[MingxSpan, MingxGeneration]]:
        parent = self._get_parent(parent_run_id)
        t = normalize_observation_type(as_type)
        if t == ObservationType.GENERATION.value:
            obs = parent.start_observation(
                name=name,
                as_type=ObservationType.GENERATION,
                input=input,
                metadata=metadata,
                model=extra.get("model"),
                model_parameters=extra.get("model_parameters"),
            )
        else:
            obs = parent.start_observation(
                name=name,
                as_type=as_type,
                input=input,
                metadata=metadata,
            )
        return obs

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        name = _get_run_name(serialized, **kwargs)
        meta = dict(metadata or {})
        if tags:
            meta["tags"] = tags
        obs_type = _get_observation_type_from_serialized(serialized, "chain", **kwargs)
        obs = self._start_obs(parent_run_id, name, obs_type, input=inputs, metadata=meta or None)
        if obs is not None:
            self._attach(run_id, obs)
            if parent_run_id is None:
                obs.update_trace(name=name, metadata=meta or None)

    def on_agent_action(
        self,
        action: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """与 Langfuse 一致：收到 agent action 时将该 run 的观测类型强制设为 agent。"""
        obs = self._runs.get(run_id)
        if obs is not None and hasattr(obs, "_otel_span"):
            obs._otel_span.set_attribute(MingxOtelSpanAttributes.OBSERVATION_TYPE, ObservationType.AGENT.value)
            if hasattr(obs, "_observation_type"):
                obs._observation_type = ObservationType.AGENT.value
            obs.update(output=action, input=kwargs.get("inputs"))

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """与 Langfuse 一致：收到 agent finish 时将该 run 的观测类型强制设为 agent。"""
        obs = self._runs.get(run_id)
        if obs is not None and hasattr(obs, "_otel_span"):
            obs._otel_span.set_attribute(MingxOtelSpanAttributes.OBSERVATION_TYPE, ObservationType.AGENT.value)
            if hasattr(obs, "_observation_type"):
                obs._observation_type = ObservationType.AGENT.value
            obs.update(output=finish, input=kwargs.get("inputs"))

    def on_chain_end(
        self,
        outputs: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(output=outputs)
            obs.end()

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(level="ERROR", status_message=str(error))
            obs.end()

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        name = _get_run_name(serialized, **kwargs)
        inv = kwargs.get("invocation_params") or {}
        model = inv.get("model_name") or inv.get("model") or serialized.get("kwargs", {}).get("model_name")
        obs = self._start_obs(
            parent_run_id,
            name,
            ObservationType.GENERATION,
            input=prompts,
            metadata=kwargs.get("metadata"),
            model=model,
            model_parameters={
                "temperature": inv.get("temperature"),
                "max_tokens": inv.get("max_tokens") or inv.get("max_completion_tokens"),
            },
        )
        if obs is not None:
            self._attach(run_id, obs)

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            output = None
            usage_details = _parse_usage_from_response(response)
            if response.generations:
                gen_list = response.generations[0] if response.generations else []
                if gen_list:
                    last_gen = gen_list[-1]
                    if isinstance(last_gen, ChatGeneration) and getattr(last_gen, "message", None):
                        output = _convert_message_to_dict(last_gen.message)
                    else:
                        texts = [g.text for g in gen_list if getattr(g, "text", None)]
                        output = texts[0] if len(texts) == 1 else texts
            obs.update(output=output, usage_details=usage_details)
            obs.end()

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(level="ERROR", status_message=str(error))
            obs.end()

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        name = _get_run_name(serialized, **kwargs)
        inv = kwargs.get("invocation_params") or {}
        model = inv.get("model_name") or inv.get("model") or serialized.get("kwargs", {}).get("model_name")
        input_val = [_convert_message_to_dict(m) for msg_list in messages for m in msg_list]
        obs = self._start_obs(
            parent_run_id,
            name,
            ObservationType.GENERATION,
            input=input_val,
            metadata=kwargs.get("metadata"),
            model=model,
            model_parameters={
                "temperature": inv.get("temperature"),
                "max_tokens": inv.get("max_tokens") or inv.get("max_completion_tokens"),
            },
        )
        if obs is not None:
            self._attach(run_id, obs)

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        name = _get_run_name(serialized, **kwargs)
        obs_type = _get_observation_type_from_serialized(serialized, "tool", **kwargs)
        obs = self._start_obs(parent_run_id, name, obs_type, input=input_str, metadata=kwargs.get("metadata"))
        if obs is not None:
            self._attach(run_id, obs)

    def on_tool_end(self, output: str, *, run_id: UUID, parent_run_id: Optional[UUID] = None, **kwargs: Any) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(output=output)
            obs.end()

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(level="ERROR", status_message=str(error))
            obs.end()

    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        name = _get_run_name(serialized, **kwargs)
        obs_type = _get_observation_type_from_serialized(serialized, "retriever", **kwargs)
        obs = self._start_obs(parent_run_id, name, obs_type, input=query, metadata=kwargs.get("metadata"))
        if obs is not None:
            self._attach(run_id, obs)

    def on_retriever_end(
        self,
        documents: Sequence[Document],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            output = [{"content": getattr(d, "page_content", str(d)), "metadata": getattr(d, "metadata", {})} for d in documents]
            obs.update(output=output)
            obs.end()

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(level="ERROR", status_message=str(error))
            obs.end()
