"""
Burrow SDK adapter for LlamaIndex.

Uses LlamaIndex's callback handler system to scan tool calls and LLM inputs
through Burrow for prompt injection detection. Automatically resolves
per-agent identity from event metadata when available.

.. deprecated::
    This adapter uses ``BaseCallbackHandler`` which is deprecated in
    LlamaIndex in favour of the instrumentation module
    (``llama_index.core.instrumentation``). A future version of this adapter
    will migrate to the new API. See
    https://docs.llamaindex.ai/en/stable/module_guides/observability/instrumentation/

Usage:
    from burrow import BurrowGuard
    from burrow.integrations.llamaindex import create_burrow_callback_handler

    guard = BurrowGuard(client_id="...", client_secret="...")
    handler = create_burrow_callback_handler(guard)

    from llama_index.core import Settings
    Settings.callback_manager.add_handler(handler)
"""

from __future__ import annotations

import json
import logging
import warnings
from typing import Any

__all__ = ["BurrowBlockedError", "create_burrow_callback_handler"]

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


class BurrowBlockedError(Exception):
    """Raised when Burrow blocks a LlamaIndex operation."""

    def __init__(self, category: str, confidence: float, context: str = ""):
        self.category = category
        self.confidence = confidence
        super().__init__(
            f"Burrow blocked{' ' + context if context else ''}: "
            f"{category} ({confidence:.0%} confidence)"
        )


def create_burrow_callback_handler(
    guard: Any,
    block_on_warn: bool = False,
) -> Any:
    """Create a LlamaIndex callback handler that scans through Burrow.

    Scans LLM inputs (user prompts), tool call arguments, and tool responses.
    Automatically resolves per-agent identity from event metadata
    (e.g. ``llamaindex:query-agent``). Falls back to ``llamaindex``.

    Args:
        guard: BurrowGuard instance.
        block_on_warn: If True, also block on "warn" verdicts.

    Returns:
        A BaseCallbackHandler to add via Settings.callback_manager.add_handler().
    """
    warnings.warn(
        "The LlamaIndex adapter uses BaseCallbackHandler which is deprecated. "
        "A future version will migrate to llama_index.core.instrumentation. "
        "See https://docs.llamaindex.ai/en/stable/module_guides/observability/instrumentation/",
        DeprecationWarning,
        stacklevel=2,
    )
    try:
        from llama_index.core.callbacks.base_handler import BaseCallbackHandler
        from llama_index.core.callbacks.schema import CBEventType
    except ImportError as exc:
        raise ImportError(
            "llama-index-core is required for the LlamaIndex adapter. "
            "Install it with: pip install burrow-sdk[llamaindex]"
        ) from exc

    class BurrowCallbackHandler(BaseCallbackHandler):
        """LlamaIndex callback handler that scans through Burrow."""

        def __init__(self, guard: Any, block_on_warn: bool):
            # llama-index-core 0.14.x renamed parameters with inverted semantics:
            # event_starts_to_trace → event_starts_to_ignore
            # event_ends_to_trace → event_ends_to_ignore
            try:
                # New API (0.14.x+): specify events to IGNORE
                all_events = list(CBEventType)
                starts_to_ignore = [
                    e for e in all_events if e not in (CBEventType.LLM, CBEventType.FUNCTION_CALL)
                ]
                ends_to_ignore = [e for e in all_events if e != CBEventType.FUNCTION_CALL]
                super().__init__(
                    event_starts_to_ignore=starts_to_ignore,
                    event_ends_to_ignore=ends_to_ignore,
                )
            except TypeError:
                # Old API (<0.14.x): specify events to TRACE
                super().__init__(
                    event_starts_to_trace=[CBEventType.LLM, CBEventType.FUNCTION_CALL],
                    event_ends_to_trace=[CBEventType.FUNCTION_CALL],
                )
            self.guard = guard
            self.block_on_warn = block_on_warn

        def on_event_start(
            self,
            event_type: Any,
            payload: dict[str, Any] | None = None,
            event_id: str = "",
            parent_id: str = "",
            **kwargs: Any,
        ) -> str:
            if payload is None:
                return event_id

            agent_name = _resolve_agent(payload, kwargs)

            if event_type == CBEventType.LLM:
                self._scan_llm_input(payload, agent_name)
            elif event_type == CBEventType.FUNCTION_CALL:
                self._scan_tool_call(payload, agent_name)

            return event_id

        def on_event_end(
            self,
            event_type: Any,
            payload: dict[str, Any] | None = None,
            event_id: str = "",
            **kwargs: Any,
        ) -> None:
            if payload is None:
                return

            if event_type == CBEventType.FUNCTION_CALL:
                agent_name = _resolve_agent(payload, kwargs)
                self._scan_tool_response(payload, agent_name)

        def start_trace(self, trace_id: str | None = None) -> None:
            pass

        def end_trace(
            self,
            trace_id: str | None = None,
            trace_map: dict[str, list[str]] | None = None,
        ) -> None:
            pass

        def _scan_llm_input(self, payload: dict[str, Any], agent_name: str) -> None:
            messages = payload.get("messages", [])
            text = ""
            if isinstance(messages, list):
                for msg in messages:
                    if isinstance(msg, dict):
                        content = msg.get("content", "")
                        if isinstance(content, str):
                            text += content + " "
                    elif hasattr(msg, "content"):
                        text += str(msg.content) + " "

            if not text.strip():
                text = str(payload.get("serialized_text", "") or payload.get("template", ""))

            text = text.strip()
            if not text:
                return

            result = self.guard.scan(
                text,
                content_type="user_prompt",
                agent=agent_name,
            )

            if result.is_blocked or (self.block_on_warn and result.is_warning):
                logger.warning(
                    "Burrow blocked LLM input for agent '%s': %s (%.0f%% confidence)",
                    agent_name,
                    result.category,
                    result.confidence * 100,
                )
                raise BurrowBlockedError(result.category, result.confidence, "LLM input")

        def _scan_tool_call(self, payload: dict[str, Any], agent_name: str) -> None:
            tool_name = payload.get("tool_name", payload.get("function_name", "unknown"))
            tool_args = payload.get("tool_args", payload.get("function_args", {}))

            text = _extract_text(tool_args)
            if not text.strip():
                return

            result = self.guard.scan(
                text,
                content_type="tool_call",
                agent=agent_name,
                tool_name=str(tool_name),
            )

            if result.is_blocked or (self.block_on_warn and result.is_warning):
                logger.warning(
                    "Burrow blocked tool call %s for agent '%s': %s (%.0f%% confidence)",
                    tool_name,
                    agent_name,
                    result.category,
                    result.confidence * 100,
                )
                raise BurrowBlockedError(
                    result.category, result.confidence, f"tool call {tool_name}"
                )

        def _scan_tool_response(self, payload: dict[str, Any], agent_name: str) -> None:
            response = payload.get("function_output", payload.get("tool_output", ""))
            text = str(response)[:4096].strip() if response else ""
            if not text:
                return

            tool_name = payload.get("tool_name", payload.get("function_name", "unknown"))

            result = self.guard.scan(
                text,
                content_type="tool_response",
                agent=agent_name,
                tool_name=str(tool_name),
            )

            if result.is_blocked or (self.block_on_warn and result.is_warning):
                logger.warning(
                    "Burrow flagged tool output from %s for agent '%s': %s (%.0f%% confidence)",
                    tool_name,
                    agent_name,
                    result.category,
                    result.confidence * 100,
                )

    return BurrowCallbackHandler(guard, block_on_warn=block_on_warn)


def _resolve_agent(payload: dict[str, Any], kwargs: dict[str, Any]) -> str:
    metadata = payload.get("metadata", kwargs.get("metadata", {}))
    if isinstance(metadata, dict):
        agent_name = metadata.get("agent_name", metadata.get("node_name", ""))
        if agent_name:
            return f"llamaindex:{agent_name}"
    return "llamaindex"


def _extract_text(tool_args: Any) -> str:
    if isinstance(tool_args, dict):
        for key in (
            "command",
            "content",
            "query",
            "text",
            "url",
            "file_path",
            "pattern",
            "input",
        ):
            if key in tool_args:
                return str(tool_args[key])
        return json.dumps(tool_args)
    if isinstance(tool_args, str):
        return tool_args
    return ""
