"""
Burrow adapter for LangChain.

Provides a callback handler that scans messages before they reach the LLM.

Usage:
    from burrow import BurrowGuard
    from burrow.integrations.langchain import create_langchain_callback

    guard = BurrowGuard(client_id="...", client_secret="...")
    callback = create_langchain_callback(guard)

    # Use with any LangChain chat model
    model = ChatOpenAI(model="gpt-4", callbacks=[callback])
"""

from __future__ import annotations

import logging
from typing import Any

from burrow import BurrowGuard, ScanResult

__all__ = ["BurrowScanError", "create_langchain_callback"]

logger = logging.getLogger("burrow.integrations.langchain")


class BurrowScanError(Exception):
    """Raised when Burrow blocks a message."""

    def __init__(self, result: ScanResult):
        self.result = result
        super().__init__(
            f"Burrow blocked message: {result.category} ({result.confidence:.0%} confidence)"
        )


def create_langchain_callback(
    guard: BurrowGuard,
    block_on_warn: bool = False,
    default_agent: str = "langchain",
):
    """
    Create a LangChain callback handler with per-agent identity.

    Automatically extracts the LangGraph node name from callback metadata,
    producing agent names like ``langchain:retriever`` or ``langchain:writer``.
    Falls back to *default_agent* when no node name is available.

    Args:
        guard: BurrowGuard instance.
        block_on_warn: If True, also block on "warn" verdicts.
        default_agent: Fallback agent name when LangGraph node is unavailable.

    Returns:
        A LangChain BaseCallbackHandler instance.
    """
    try:
        from langchain_core.callbacks import BaseCallbackHandler
    except ImportError as exc:
        raise ImportError(
            "langchain-core is required for the LangChain adapter. "
            "Install it with: pip install burrow-sdk[langchain]"
        ) from exc

    def _resolve_agent(kwargs: dict) -> str:
        node = kwargs.get("metadata", {}).get("langgraph_node")
        if node:
            return f"langchain:{node}"
        return default_agent

    class BurrowCallbackHandler(BaseCallbackHandler):
        """LangChain callback with per-agent identity via LangGraph metadata."""

        name = "burrow_guard"

        def on_llm_start(
            self,
            serialized: dict[str, Any],
            prompts: list[str],
            **kwargs: Any,
        ) -> None:
            agent_name = _resolve_agent(kwargs)
            for prompt_text in prompts:
                if not prompt_text.strip():
                    continue
                result = guard.scan(
                    prompt_text,
                    content_type="user_prompt",
                    agent=agent_name,
                )
                if result.is_blocked or (block_on_warn and result.is_warning):
                    raise BurrowScanError(result)

        def on_chat_model_start(
            self,
            serialized: dict[str, Any],
            messages: list[list[Any]],
            **kwargs: Any,
        ) -> None:
            agent_name = _resolve_agent(kwargs)
            for message_batch in messages:
                for msg in message_batch:
                    content = getattr(msg, "content", "")
                    if isinstance(content, str) and content.strip():
                        msg_type = getattr(msg, "type", "human")
                        content_type = "tool_response" if msg_type == "tool" else "user_prompt"
                        result = guard.scan(
                            content,
                            content_type=content_type,
                            agent=agent_name,
                        )
                        if result.is_blocked or (block_on_warn and result.is_warning):
                            raise BurrowScanError(result)

        def on_tool_end(self, output: str, **kwargs: Any) -> None:
            agent_name = _resolve_agent(kwargs)
            if isinstance(output, str) and output.strip():
                result = guard.scan(
                    output,
                    content_type="tool_response",
                    agent=agent_name,
                    tool_name=kwargs.get("name"),
                )
                if result.is_blocked or (block_on_warn and result.is_warning):
                    raise BurrowScanError(result)

    return BurrowCallbackHandler()
