"""
LangChain callback handler for Raindrop observability.

Captures LLM calls, tool usage, chains, and retrievers via LangChain's
callback system and ships them to the Raindrop API.

Usage:
    from raindrop_langchain import create_raindrop_langchain

    raindrop = create_raindrop_langchain(api_key="rk_...")
    model = ChatOpenAI()
    result = model.invoke("Hello", config={"callbacks": [raindrop["handler"]]})
    raindrop["flush"]()
"""

from __future__ import annotations

import json
import uuid
from typing import Any, Dict, List, Optional, Sequence

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

import raindrop.analytics as raindrop

_LANGGRAPH_INTERNAL_NAMES = frozenset({
    "LangGraph",
    "__start__",
    "__end__",
    "ChannelWrite",
    "ChannelRead",
})


def _is_langgraph_internal(name: str) -> bool:
    if name in _LANGGRAPH_INTERNAL_NAMES:
        return True
    if name.startswith("Branch:"):
        return True
    if name.startswith("__") and name.endswith("__"):
        return True
    return False


class RaindropCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler that ships events to Raindrop."""

    def __init__(
        self,
        api_key: str,
        user_id: Optional[str] = None,
        convo_id: Optional[str] = None,
        trace_chains: bool = True,
        trace_retrievers: bool = True,
        filter_langgraph_internals: bool = True,
    ):
        super().__init__()
        self._user_id = user_id
        self._convo_id = convo_id
        self._trace_chains = trace_chains
        self._trace_retrievers = trace_retrievers
        self._filter_langgraph = filter_langgraph_internals

        raindrop.init(api_key=api_key)

        self._spans: Dict[str, Any] = {}
        self._root_run_ids: Dict[str, str] = {}
        self._event_ids: Dict[str, str] = {}

    def _get_event_id(self, run_id: uuid.UUID, parent_run_id: Optional[uuid.UUID] = None) -> str:
        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else None
        root_id = self._root_run_ids.get(pid, pid) if pid else rid
        self._root_run_ids[rid] = root_id
        if root_id not in self._event_ids:
            self._event_ids[root_id] = str(uuid.uuid4())
        return self._event_ids[root_id]

    def _cleanup(self, run_id: uuid.UUID) -> None:
        rid = str(run_id)
        self._spans.pop(rid, None)
        self._spans.pop(f"{run_id}:agent_action", None)
        root_id = self._root_run_ids.pop(rid, rid)
        if root_id == rid:
            self._event_ids.pop(root_id, None)

    def _finalize_if_root(
        self,
        run_id: uuid.UUID,
        output: Optional[str] = None,
        model: Optional[str] = None,
        properties: Optional[Dict[str, Any]] = None,
    ) -> None:
        rid = str(run_id)
        root_id = self._root_run_ids.get(rid, rid)
        event_id = self._event_ids.get(root_id)
        if not event_id:
            return

        if root_id == rid:
            raindrop.track_ai(
                user_id=self._user_id or "unknown",
                event="ai_generation",
                event_id=event_id,
                output=output,
                model=model,
                properties=properties,
                convo_id=self._convo_id,
            )

    @staticmethod
    def _build_properties(
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        if not tags and not metadata:
            return None
        props: Dict[str, Any] = {}
        if tags:
            props["langchain.tags"] = tags
        if metadata:
            props.update(metadata)
        return props

    # -------------------------------------------------------------------------
    # LLM events
    # -------------------------------------------------------------------------

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        try:
            if self._filter_langgraph and str(run_id) in self._spans:
                return
            event_id = self._get_event_id(run_id, parent_run_id)
            self._spans[str(run_id)] = {
                "event_id": event_id,
                "type": "llm",
                "input": "\n".join(prompts),
                "properties": self._build_properties(tags, metadata),
            }
        except Exception:
            pass

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        try:
            if self._filter_langgraph and str(run_id) in self._spans:
                return
            event_id = self._get_event_id(run_id, parent_run_id)
            self._spans[str(run_id)] = {
                "event_id": event_id,
                "type": "chat",
                "properties": self._build_properties(tags, metadata),
            }

            all_msgs = [msg for msg_list in messages for msg in msg_list]
            user_msgs = [m for m in all_msgs if getattr(m, "type", "unknown") in ("human", "user")]
            last_msg = user_msgs[-1] if user_msgs else (all_msgs[-1] if all_msgs else None)
            input_text = None
            if last_msg:
                content = last_msg.content if hasattr(last_msg, "content") else str(last_msg)
                if isinstance(content, list):
                    try:
                        content = json.dumps(content)
                    except Exception:
                        content = str(content)
                input_text = content
            self._spans[str(run_id)]["input"] = input_text
        except Exception:
            pass

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            model, output_text, prompt_tokens, completion_tokens = _extract_llm_metadata(response)

            span_data = self._spans.get(str(run_id), {})
            input_text = span_data.get("input")
            span_props = span_data.get("properties") or {}

            properties: Dict[str, Any] = {**span_props}
            if prompt_tokens is not None:
                properties["ai.usage.prompt_tokens"] = prompt_tokens
            if completion_tokens is not None:
                properties["ai.usage.completion_tokens"] = completion_tokens

            rid = str(run_id)
            root_id = self._root_run_ids.get(rid, rid)
            event_id = self._event_ids.get(root_id)

            if event_id:
                raindrop.track_ai(
                    user_id=self._user_id or "unknown",
                    event="ai_generation",
                    event_id=event_id if root_id == rid else str(uuid.uuid4()),
                    input=input_text,
                    output=output_text,
                    model=model,
                    properties=properties if properties else None,
                    convo_id=self._convo_id,
                )
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            # NOTE: Input captured in on_llm_start is in _spans, not available
            # to _finalize_if_root, so the finalized event will have input=None on error paths.
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    # -------------------------------------------------------------------------
    # Chain events
    # -------------------------------------------------------------------------

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        name: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        try:
            if self._filter_langgraph:
                chain_name = name or _get_name_from_serialized(serialized)
                if _is_langgraph_internal(chain_name):
                    return
            self._get_event_id(run_id, parent_run_id)
            if not self._trace_chains:
                return
            self._spans[str(run_id)] = {"type": "chain"}
        except Exception:
            pass

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    # -------------------------------------------------------------------------
    # Tool events
    # -------------------------------------------------------------------------

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        try:
            event_id = self._get_event_id(run_id, parent_run_id)
            tool_name = serialized.get("name") or (
                serialized.get("id", [None])[-1] if serialized.get("id") else None
            )
            self._spans[str(run_id)] = {
                "event_id": event_id,
                "type": "tool",
                "name": tool_name,
            }
        except Exception:
            pass

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    # -------------------------------------------------------------------------
    # Retriever events
    # -------------------------------------------------------------------------

    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._get_event_id(run_id, parent_run_id)
            if not self._trace_retrievers:
                return
            self._spans[str(run_id)] = {"type": "retriever"}
        except Exception:
            pass

    def on_retriever_end(
        self,
        documents: Sequence[Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._finalize_if_root(run_id)
        except Exception:
            pass
        finally:
            self._cleanup(run_id)

    # -------------------------------------------------------------------------
    # Agent events
    # -------------------------------------------------------------------------

    def on_agent_action(
        self,
        action: Any,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            self._get_event_id(run_id, parent_run_id)
            tool = getattr(action, "tool", None) or "unknown"
            tool_input = getattr(action, "tool_input", "")
            if not isinstance(tool_input, str):
                try:
                    tool_input = json.dumps(tool_input)
                except Exception:
                    tool_input = str(tool_input)
            span_key = f"{run_id}:agent_action"
            self._spans[span_key] = {
                "type": "agent_action",
                "tool": tool,
                "tool_input": tool_input,
            }
        except Exception:
            pass

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            span_key = f"{run_id}:agent_action"
            self._spans.pop(span_key, None)
        except Exception:
            pass


def create_raindrop_langchain(
    api_key: str,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
    trace_chains: bool = True,
    trace_retrievers: bool = True,
    filter_langgraph_internals: bool = True,
) -> dict:
    """
    Create a Raindrop-instrumented LangChain callback handler.

    Args:
        api_key: Raindrop API key (rk_...).
        user_id: Associate all events with this user.
        convo_id: Conversation ID to group related events.
        trace_chains: Track chain execution (default: True).
        trace_retrievers: Track retriever calls (default: True).
        filter_langgraph_internals: Filter LangGraph-internal chain events
            and deduplicate LLM callbacks (default: True).

    Returns:
        Dict with 'handler', 'flush', and 'shutdown'.

    Usage:
        raindrop = create_raindrop_langchain(api_key="rk_...")
        model = ChatOpenAI()
        result = model.invoke("Hello", config={"callbacks": [raindrop["handler"]]})
        raindrop["flush"]()
    """
    handler = RaindropCallbackHandler(
        api_key=api_key,
        user_id=user_id,
        convo_id=convo_id,
        trace_chains=trace_chains,
        trace_retrievers=trace_retrievers,
        filter_langgraph_internals=filter_langgraph_internals,
    )

    return {
        "handler": handler,
        "flush": raindrop.flush,
        "shutdown": raindrop.shutdown,
    }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _get_name_from_serialized(serialized: Dict[str, Any]) -> str:
    """Extract a human-readable name from a LangChain serialized object."""
    if serialized.get("name"):
        return serialized["name"]
    ids = serialized.get("id")
    if ids and isinstance(ids, list) and len(ids) > 0:
        return ids[-1]
    return "unknown"


def _extract_llm_metadata(response: LLMResult):
    """Extract model name, output text, and token counts from an LLMResult."""
    model = None
    output_text = None
    prompt_tokens = None
    completion_tokens = None

    if response.generations and response.generations[0]:
        gen = response.generations[0][0]
        output_text = gen.text

        msg = getattr(gen, "message", None)
        if msg:
            resp_meta = getattr(msg, "response_metadata", {}) or {}
            model = resp_meta.get("model_name") or resp_meta.get("model")

            usage_meta = getattr(msg, "usage_metadata", None)
            if usage_meta:
                if isinstance(usage_meta, dict):
                    prompt_tokens = usage_meta.get("input_tokens")
                    completion_tokens = usage_meta.get("output_tokens")
                else:
                    prompt_tokens = getattr(usage_meta, "input_tokens", None)
                    completion_tokens = getattr(usage_meta, "output_tokens", None)

    llm_output = response.llm_output or {}
    if not model:
        model = llm_output.get("model_name") or llm_output.get("model")

    token_usage = llm_output.get("token_usage") or llm_output.get("tokenUsage") or {}
    if prompt_tokens is None:
        pt = token_usage.get("prompt_tokens")
        prompt_tokens = pt if pt is not None else token_usage.get("promptTokens")
    if completion_tokens is None:
        ct = token_usage.get("completion_tokens")
        completion_tokens = ct if ct is not None else token_usage.get("completionTokens")

    return model, output_text, prompt_tokens, completion_tokens
