"""
LangChain integration — provides a callback handler that auto-traces all
LLM / ChatModel / Chain / Tool calls through the Syntropy platform.

Usage with a ChatModel:

    from langchain_openai import ChatOpenAI
    from syntropy import Syntropy
    from syntropy.integrations.langchain import SyntropyCallbackHandler

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")
    handler = SyntropyCallbackHandler(syn)

    model = ChatOpenAI(model="gpt-4o", callbacks=[handler])
    result = model.invoke("Hello!")

Usage with a chain:

    from langchain_core.runnables import RunnableSequence

    chain = prompt | model | parser
    result = chain.invoke({"input": "Hello!"}, config={"callbacks": [handler]})
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Union
from uuid import UUID

from syntropy.client import Syntropy
from syntropy.trace_context import trace_scope

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.messages import BaseMessage
    from langchain_core.outputs import ChatResult, LLMResult

    _HAS_LANGCHAIN = True
except ImportError:
    _HAS_LANGCHAIN = False

    # Provide a minimal base so the class definition still works
    # when langchain-core is not installed (duck-typing approach).
    class BaseCallbackHandler:  # type: ignore[no-redef]
        """Stub — real BaseCallbackHandler comes from langchain-core."""

        pass

    BaseMessage = Any  # type: ignore[assignment,misc]
    LLMResult = Any  # type: ignore[assignment,misc]
    ChatResult = Any  # type: ignore[assignment,misc]


# ================================================================
# Internal state
# ================================================================


@dataclass
class _RunState:
    """Tracks in-flight LLM / chain / tool runs."""

    start_time: float
    prompt: str
    model: str
    agent_id: str
    metadata: Dict[str, Any] = field(default_factory=dict)


# ================================================================
# Callback handler
# ================================================================


class SyntropyCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler that automatically sends traces to the
    Syntropy ingestion service for every LLM, Chat Model, Chain, and
    Tool invocation.

    The handler is fully compatible with LangChain's ``BaseCallbackHandler``
    interface. If ``langchain-core`` is not installed it still works via
    duck-typing — LangChain dispatches callbacks by method name, not
    ``isinstance`` checks.
    """

    name: str = "SyntropyCallbackHandler"

    def __init__(
        self,
        syntropy: Syntropy,
        *,
        agent_id: str = "",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        if _HAS_LANGCHAIN:
            super().__init__()
        self._syntropy = syntropy
        self._agent_id = agent_id or syntropy._config.agent_id
        self._default_meta: Dict[str, Any] = metadata or {}
        self._runs: Dict[str, _RunState] = {}
        self._scope_depth: int = 0

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

    def _enter_scope(self) -> None:
        """Mark that a LangChain LLM call is in progress.

        When ``_scope_depth > 0``, any Syntropy provider wrappers that
        check ``is_inside_trace()`` will skip their own trace dispatch so
        the call is not double-counted.
        """
        from syntropy.trace_context import _active_trace
        if self._scope_depth == 0:
            _active_trace.set("langchain")
        self._scope_depth += 1

    def _exit_scope(self) -> None:
        """Release the dedup scope opened by ``_enter_scope``."""
        from syntropy.trace_context import _active_trace
        self._scope_depth = max(0, self._scope_depth - 1)
        if self._scope_depth == 0:
            _active_trace.set(None)

    def _extract_model(
        self,
        serialized: Dict[str, Any],
        kwargs: Optional[Dict[str, Any]] = None,
    ) -> str:
        """Best-effort model name extraction from serialized + kwargs."""
        if kwargs:
            inv = kwargs.get("invocation_params", {})
            for key in ("model", "model_name", "model_id"):
                if inv.get(key):
                    return str(inv[key])
        # Fall back to the serialized class hierarchy id
        ids = serialized.get("id", [])
        return str(ids[-1]) if ids else "unknown"

    @staticmethod
    def _msg_to_str(msg: Any) -> str:
        if isinstance(msg, str):
            return msg
        if hasattr(msg, "content"):
            c = msg.content
            return c if isinstance(c, str) else json.dumps(c)
        return str(msg)

    # ----------------------------------------------------------------
    # LLM callbacks
    # ----------------------------------------------------------------

    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,
    ) -> None:
        # Enter trace scope so nested provider wrappers skip their own trace
        self._enter_scope()

        self._runs[str(run_id)] = _RunState(
            start_time=time.perf_counter(),
            prompt="\n".join(prompts),
            model=self._extract_model(serialized, kwargs),
            agent_id=self._agent_id,
            metadata={**self._default_meta, **(metadata or {})},
        )

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        # Release the dedup scope
        self._exit_scope()

        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        # Collect generation texts
        texts: List[str] = []
        for gen_list in getattr(response, "generations", []):
            for g in gen_list:
                t = getattr(g, "text", "")
                if t:
                    texts.append(t)
        output = "\n".join(texts)

        # Token usage
        llm_output = getattr(response, "llm_output", {}) or {}
        usage = llm_output.get("token_usage", llm_output)
        tokens_in = usage.get("prompt_tokens", 0)
        tokens_out = usage.get("completion_tokens", 0)

        self._syntropy.trace(
            prompt=run.prompt,
            response=output,
            model=run.model,
            agent_id=run.agent_id or None,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "llm",
            },
        )

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        # Release the dedup scope
        self._exit_scope()

        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        self._syntropy.trace(
            prompt=run.prompt,
            response="",
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            status="error",
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "llm",
                "error": str(error),
            },
        )

    # ----------------------------------------------------------------
    # Chat model callbacks
    # ----------------------------------------------------------------

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        # Enter trace scope so nested provider wrappers skip their own trace
        self._enter_scope()

        flat = [m for batch in messages for m in batch]
        last = flat[-1] if flat else None
        prompt = self._msg_to_str(last) if last else ""

        self._runs[str(run_id)] = _RunState(
            start_time=time.perf_counter(),
            prompt=prompt,
            model=self._extract_model(serialized, kwargs),
            agent_id=self._agent_id,
            metadata={
                **self._default_meta,
                **(metadata or {}),
                "message_count": len(flat),
            },
        )

    # on_llm_end / on_llm_error handle chat model completion events.

    # ----------------------------------------------------------------
    # Chain callbacks
    # ----------------------------------------------------------------

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        prompt = inputs if isinstance(inputs, str) else json.dumps(inputs, default=str)

        self._runs[str(run_id)] = _RunState(
            start_time=time.perf_counter(),
            prompt=prompt,
            model=serialized.get("id", ["chain"])[-1],
            agent_id=self._agent_id,
            metadata={**self._default_meta, "integration": "chain"},
        )

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)
        response = outputs if isinstance(outputs, str) else json.dumps(outputs, default=str)

        self._syntropy.trace(
            prompt=run.prompt,
            response=response,
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "chain",
            },
        )

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        self._syntropy.trace(
            prompt=run.prompt,
            response="",
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            status="error",
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "chain",
                "error": str(error),
            },
        )

    # ----------------------------------------------------------------
    # Tool callbacks
    # ----------------------------------------------------------------

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        tool_name = serialized.get("name") or serialized.get("id", ["unknown"])[-1]

        self._runs[str(run_id)] = _RunState(
            start_time=time.perf_counter(),
            prompt=input_str,
            model=str(tool_name),
            agent_id=self._agent_id,
            metadata={
                **self._default_meta,
                "integration": "tool",
                "tool_name": str(tool_name),
            },
        )

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        self._syntropy.trace(
            prompt=run.prompt,
            response=output if isinstance(output, str) else str(output),
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            metadata={
                **run.metadata,
                "provider": "langchain",
            },
        )

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        self._syntropy.trace(
            prompt=run.prompt,
            response="",
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            status="error",
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "tool",
                "error": str(error),
            },
        )

    # ----------------------------------------------------------------
    # Retriever callbacks (for RAG tracing)
    # ----------------------------------------------------------------

    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        self._runs[str(run_id)] = _RunState(
            start_time=time.perf_counter(),
            prompt=query,
            model="retriever",
            agent_id=self._agent_id,
            metadata={
                **self._default_meta,
                "integration": "retriever",
            },
        )

    def on_retriever_end(
        self,
        documents: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        # Summarize retrieved documents
        doc_count = len(documents) if documents else 0
        doc_texts = []
        for doc in (documents or [])[:5]:
            content = getattr(doc, "page_content", str(doc))
            doc_texts.append(content[:200])

        self._syntropy.trace(
            prompt=run.prompt,
            response=json.dumps(
                {"document_count": doc_count, "preview": doc_texts},
                default=str,
            ),
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "retriever",
                "document_count": doc_count,
            },
        )

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run = self._runs.pop(str(run_id), None)
        if not run:
            return

        elapsed_ms = int((time.perf_counter() - run.start_time) * 1000)

        self._syntropy.trace(
            prompt=run.prompt,
            response="",
            model=run.model,
            agent_id=run.agent_id or None,
            latency_ms=elapsed_ms,
            status="error",
            metadata={
                **run.metadata,
                "provider": "langchain",
                "integration": "retriever",
                "error": str(error),
            },
        )
