"""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 BaseMessage
    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 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 _message_to_dict(msg: BaseMessage) -> Dict[str, Any]:
    if hasattr(msg, "type") and hasattr(msg, "content"):
        return {"role": getattr(msg, "type", "message"), "content": getattr(msg, "content", "")}
    return {"content": str(msg)}


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  # token for context.detach(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:
                context.detach(token)
            except Exception:
                pass
        return obs

    def _start_obs(
        self,
        parent_run_id: Optional[UUID],
        name: str,
        as_type: str = "span",
        input: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **extra: Any,
    ) -> Optional[Union[MingxSpan, MingxGeneration]]:
        parent = self._get_parent(parent_run_id)
        if as_type == "generation":
            obs = parent.start_observation(
                name=name,
                as_type="generation",
                input=input,
                metadata=metadata,
                model=extra.get("model"),
                model_parameters=extra.get("model_parameters"),
                **{k: v for k, v in extra.items() if k in ("model", "model_parameters")},
            )
        else:
            obs = parent.start_observation(
                name=name,
                as_type="span",
                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 = self._start_obs(parent_run_id, name, "span", 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, input=inputs, metadata=meta or None)

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

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

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        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)
        inv = kwargs.get("invocation_params") or {}
        model = inv.get("model_name") or inv.get("model") or serialized.get("kwargs", {}).get("model_name")
        meta = dict(metadata or {})
        if tags:
            meta["tags"] = tags
        parent = self._get_parent(parent_run_id)
        obs = parent.start_observation(
            name=name,
            as_type="generation",
            input=prompts[0] if len(prompts) == 1 else prompts,
            metadata=meta or None,
            model=model,
            model_parameters={
                k: v
                for k, v in {
                    "temperature": inv.get("temperature"),
                    "max_tokens": inv.get("max_tokens") or inv.get("max_completion_tokens"),
                }.items()
                if v is not None
            } or None,
        )
        if obs is not None:
            self._attach(run_id, obs)

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        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)
        inv = kwargs.get("invocation_params") or {}
        model = inv.get("model_name") or inv.get("model")
        flat = [m for row in messages for m in row]
        input_list = [_message_to_dict(m) for m in flat]
        meta = dict(metadata or {})
        if tags:
            meta["tags"] = tags
        parent = self._get_parent(parent_run_id)
        obs = parent.start_observation(
            name=name,
            as_type="generation",
            input=input_list,
            metadata=meta or None,
            model=model,
            model_parameters={
                k: v
                for k, v in {
                    "temperature": inv.get("temperature"),
                    "max_tokens": inv.get("max_tokens") or inv.get("max_completion_tokens"),
                }.items()
                if v is not None
            } or None,
        )
        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,
    ) -> Any:
        obs = self._detach(run_id)
        if obs is None:
            return
        generations = response.generations
        output = None
        usage_details = None
        if generations and generations[0]:
            gen0 = generations[0][0]
            if isinstance(gen0, ChatGeneration) and hasattr(gen0, "message"):
                msg = gen0.message
                output = getattr(msg, "content", str(msg))
            else:
                output = getattr(gen0, "text", str(gen0))
            if hasattr(gen0, "generation_info") and gen0.generation_info:
                ui = gen0.generation_info
                usage_details = {}
                if ui.get("prompt_tokens") is not None:
                    usage_details["prompt_tokens"] = ui["prompt_tokens"]
                if ui.get("completion_tokens") is not None:
                    usage_details["completion_tokens"] = ui["completion_tokens"]
                if ui.get("total_tokens") is not None:
                    usage_details["total_tokens"] = ui["total_tokens"]
        if hasattr(response, "llm_output") and response.llm_output and usage_details is None:
            lo = response.llm_output
            if "token_usage" in lo:
                usage_details = dict(lo["token_usage"])
            elif "usage" in lo:
                usage_details = dict(lo["usage"])
        if isinstance(obs, MingxGeneration):
            obs.update(output=output, usage_details=usage_details, input=kwargs.get("inputs"))
        else:
            obs.update(output=output, input=kwargs.get("inputs"))
        obs.end()

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

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        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 = self._start_obs(parent_run_id, name, "span", input=input_str, metadata=meta or None)
        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) -> Any:
        obs = self._detach(run_id)
        if obs is not None:
            obs.update(output=output, input=kwargs.get("inputs")).end()

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

    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        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 = self._start_obs(parent_run_id, name, "span", input=query, metadata=meta or None)
        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,
    ) -> Any:
        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, input=kwargs.get("inputs")).end()

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