"""
Burrow adapter for LiteLLM.

Provides two integration modes:

1. SDK mode (BurrowCallback): Scans prompts and responses when using litellm
   as a Python library. Pre-call scanning raises to block; post-call scanning
   audits responses.

2. Proxy mode (BurrowGuardrail): Scans prompts and responses when running
   LiteLLM as a proxy server. Can block requests before they reach the LLM.

SDK mode usage:
    from burrow import BurrowGuard
    from burrow.integrations.litellm import create_burrow_callback

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

    import litellm
    litellm.callbacks = [callback]

    response = litellm.completion(model="gpt-4", messages=[...])

Proxy mode usage (config.yaml):
    guardrails:
      - guardrail_name: "burrow-firewall"
        litellm_params:
          guardrail: burrow.integrations.litellm.BurrowGuardrail
          mode: "pre_call"
"""

from __future__ import annotations

import logging
import os

from burrow import BurrowGuard, ScanResult

__all__ = ["BurrowBlockedError", "BurrowGuardrail", "create_burrow_callback"]

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


class BurrowBlockedError(Exception):
    """Raised when Burrow blocks an LLM request."""

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


def _extract_user_message(messages: list) -> str:
    """Extract the last user message from the messages list."""
    for msg in reversed(messages):
        if isinstance(msg, dict) and msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                return " ".join(
                    p.get("text", "")
                    for p in content
                    if isinstance(p, dict) and p.get("type") == "text"
                )
    return ""


def _should_block(result: ScanResult, block_on_warn: bool) -> bool:
    return result.is_blocked or (block_on_warn and result.is_warning)


# ---------------------------------------------------------------------------
# SDK mode: CustomLogger callback
# ---------------------------------------------------------------------------


def create_burrow_callback(
    guard: BurrowGuard,
    agent: str = "litellm",
    scan_responses: bool = True,
    block_on_warn: bool = False,
):
    """
    Create a LiteLLM callback that scans prompts and responses with Burrow.

    Works with LiteLLM's callback system (CustomLogger). Register it with:
        litellm.callbacks = [callback]

    Note: In SDK mode, pre-call blocking works by raising from log_pre_api_call.
    LiteLLM swallows these exceptions in some code paths. For guaranteed blocking,
    call guard.scan() before litellm.completion() directly.

    Args:
        guard: BurrowGuard instance.
        agent: Agent name for scan context.
        scan_responses: If True, also scan LLM responses.
        block_on_warn: If True, also block on "warn" verdicts.

    Returns:
        A LiteLLM CustomLogger instance.
    """
    try:
        from litellm.integrations.custom_logger import CustomLogger
    except ImportError as exc:
        raise ImportError(
            "litellm is required for the LiteLLM adapter. "
            "Install it with: pip install burrow-sdk[litellm]"
        ) from exc

    class BurrowCallback(CustomLogger):
        """LiteLLM callback that scans LLM inputs/outputs with Burrow."""

        def log_pre_api_call(self, model, messages, kwargs):
            """Scan the user prompt before the LLM call."""
            text = _extract_user_message(messages)
            if not text.strip():
                return

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

            if _should_block(result, block_on_warn):
                logger.warning(
                    "burrow blocked pre-call: model=%s category=%s confidence=%.2f request_id=%s",
                    model,
                    result.category,
                    result.confidence,
                    result.request_id,
                )
                raise BurrowBlockedError(result)

            if result.is_warning:
                logger.info(
                    "burrow warn pre-call: model=%s category=%s confidence=%.2f",
                    model,
                    result.category,
                    result.confidence,
                )

        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
            """Scan the LLM response after a successful call."""
            if not scan_responses:
                return

            try:
                content = response_obj.choices[0].message.content
            except (AttributeError, IndexError):
                return

            if not content or not content.strip():
                return

            result = await guard.scan_async(
                content,
                content_type="tool_response",
                agent=agent,
            )

            if _should_block(result, block_on_warn):
                logger.warning(
                    "burrow blocked response: model=%s category=%s confidence=%.2f request_id=%s",
                    kwargs.get("model", "unknown"),
                    result.category,
                    result.confidence,
                    result.request_id,
                )

            # Scan tool calls if present
            try:
                tool_calls = response_obj.choices[0].message.tool_calls
            except (AttributeError, IndexError):
                tool_calls = None

            if tool_calls:
                for tc in tool_calls:
                    fn = getattr(tc, "function", None)
                    if fn and hasattr(fn, "arguments"):
                        tc_result = await guard.scan_async(
                            fn.arguments,
                            content_type="tool_call",
                            agent=agent,
                            tool_name=fn.name if hasattr(fn, "name") else None,
                        )
                        if _should_block(tc_result, block_on_warn):
                            logger.warning(
                                "burrow blocked tool_call: tool=%s category=%s confidence=%.2f",
                                getattr(fn, "name", "unknown"),
                                tc_result.category,
                                tc_result.confidence,
                            )

    return BurrowCallback()


# ---------------------------------------------------------------------------
# Proxy mode: CustomGuardrail (can block requests)
# ---------------------------------------------------------------------------


class BurrowGuardrail:
    """
    LiteLLM Proxy guardrail that scans prompts and responses with Burrow.

    This class extends CustomGuardrail and can block requests before they
    reach the LLM. Only works when LiteLLM is running as a proxy server.

    Register in LiteLLM proxy config.yaml:
        guardrails:
          - guardrail_name: "burrow-firewall"
            litellm_params:
              guardrail: burrow.integrations.litellm.BurrowGuardrail
              mode: "pre_call"

    Or set environment variables:
        BURROW_CLIENT_ID=...
        BURROW_CLIENT_SECRET=...
        BURROW_API_URL=https://api.burrow.run  (optional)
    """

    def __new__(cls, *args, **kwargs):
        try:
            from litellm.integrations.custom_guardrail import CustomGuardrail as _Base
        except ImportError as exc:
            raise ImportError(
                "litellm is required for the LiteLLM guardrail. "
                "Install it with: pip install burrow-sdk[litellm]"
            ) from exc

        guard = BurrowGuard(
            client_id=kwargs.pop("client_id", None),
            client_secret=kwargs.pop("client_secret", None),
            api_url=kwargs.pop(
                "api_url", os.environ.get("BURROW_API_URL", "https://api.burrow.run")
            ),
            fail_open=kwargs.pop("fail_open", True),
        )
        agent_name = kwargs.pop("agent", "litellm-proxy")
        block_on_warn_val = kwargs.pop("block_on_warn", False)

        class _BurrowGuardrailImpl(_Base):
            def __init__(self, **kw):
                super().__init__(
                    guardrail_name=kw.pop("guardrail_name", "burrow-firewall"),
                    event_hook=kw.pop("event_hook", ["pre_call", "post_call"]),
                    default_on=kw.pop("default_on", True),
                    **kw,
                )
                self._guard = guard
                self._agent = agent_name
                self._block_on_warn = block_on_warn_val

            async def async_pre_call_hook(
                self,
                user_api_key_dict,
                cache,
                data: dict,
                call_type: str,
            ):
                """Scan prompt before LLM call. Raise to block."""
                messages = data.get("messages", [])
                text = _extract_user_message(messages)
                if not text.strip():
                    return data

                result = await self._guard.scan_async(
                    text, content_type="user_prompt", agent=self._agent
                )

                if _should_block(result, self._block_on_warn):
                    logger.warning(
                        "burrow blocked proxy request: call_type=%s category=%s "
                        "confidence=%.2f request_id=%s",
                        call_type,
                        result.category,
                        result.confidence,
                        result.request_id,
                    )
                    self.raise_passthrough_exception(
                        violation_message=(
                            f"Burrow: Prompt injection detected "
                            f"({result.category}, {result.confidence:.0%} confidence)"
                        ),
                        request_data=data,
                        detection_info={
                            "action": result.action,
                            "confidence": result.confidence,
                            "category": result.category,
                            "request_id": result.request_id,
                        },
                    )

                if result.is_warning:
                    logger.info(
                        "burrow warn proxy request: call_type=%s category=%s confidence=%.2f",
                        call_type,
                        result.category,
                        result.confidence,
                    )

                return data

            async def async_post_call_success_hook(
                self,
                data: dict,
                user_api_key_dict,
                response,
            ):
                """Scan LLM response for injection in output."""
                try:
                    content = response.choices[0].message.content
                except (AttributeError, IndexError):
                    return response

                if not content or not content.strip():
                    return response

                result = await self._guard.scan_async(
                    content,
                    content_type="tool_response",
                    agent=self._agent,
                )

                if _should_block(result, self._block_on_warn):
                    logger.warning(
                        "burrow blocked proxy response: category=%s confidence=%.2f request_id=%s",
                        result.category,
                        result.confidence,
                        result.request_id,
                    )

                return response

        return _BurrowGuardrailImpl(**kwargs)
