"""
Burrow adapter for Google Vertex AI.

Wraps Vertex AI's GenerativeModel to scan prompts before they reach the
model and optionally audit responses.

Usage:
    from burrow import BurrowGuard
    from burrow.integrations.vertex import create_guarded_model

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

    model = GenerativeModel("gemini-1.5-pro")
    safe_model = create_guarded_model(model, guard)

    response = safe_model.generate_content("Summarize this report...")
"""

from __future__ import annotations

import logging

from burrow import BurrowGuard, ScanResult

__all__ = ["BurrowBlockedError", "create_guarded_model"]

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


class BurrowBlockedError(Exception):
    """Raised when Burrow blocks a Vertex AI request."""

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


def create_guarded_model(
    model,
    guard: BurrowGuard,
    agent: str = "vertex-ai",
    block_on_warn: bool = False,
    scan_responses: bool = False,
):
    """
    Wrap a Vertex AI GenerativeModel with Burrow prompt scanning.

    Intercepts generate_content() and generate_content_async() to scan
    the prompt before it reaches the model. If Burrow blocks the prompt,
    a BurrowBlockedError is raised and the model is never called.

    Args:
        model: A vertexai.generative_models.GenerativeModel instance.
        guard: BurrowGuard instance for scanning.
        agent: Agent name for scan context.
        block_on_warn: If True, also block on "warn" verdicts.
        scan_responses: If True, scan model responses for injection.

    Returns:
        A wrapped model with the same interface.
    """
    original_generate = model.generate_content
    original_generate_async = getattr(model, "generate_content_async", None)

    def _extract_text(contents) -> str:
        """Extract text from Vertex AI content formats."""
        if isinstance(contents, str):
            return contents

        if isinstance(contents, list):
            parts = []
            for item in contents:
                if isinstance(item, str):
                    parts.append(item)
                elif hasattr(item, "text"):
                    parts.append(item.text)
                elif hasattr(item, "parts"):
                    for part in item.parts:
                        if hasattr(part, "text"):
                            parts.append(part.text)
            return " ".join(parts)

        if hasattr(contents, "text"):
            return contents.text
        if hasattr(contents, "parts"):
            return " ".join(p.text for p in contents.parts if hasattr(p, "text"))

        return str(contents)

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

    def guarded_generate_content(contents, *args, **kwargs):
        text = _extract_text(contents)
        if text.strip():
            result = guard.scan(text, content_type="user_prompt", agent=agent)
            if _should_block(result):
                logger.warning(
                    "burrow blocked vertex request: category=%s confidence=%.2f request_id=%s",
                    result.category,
                    result.confidence,
                    result.request_id,
                )
                raise BurrowBlockedError(result)
            if result.is_warning:
                logger.info(
                    "burrow warn vertex request: category=%s confidence=%.2f",
                    result.category,
                    result.confidence,
                )

        response = original_generate(contents, *args, **kwargs)

        if scan_responses:
            try:
                resp_text = response.text
                if resp_text and resp_text.strip():
                    resp_result = guard.scan(
                        resp_text,
                        content_type="tool_response",
                        agent=agent,
                    )
                    if _should_block(resp_result):
                        logger.warning(
                            "burrow blocked vertex response: category=%s confidence=%.2f",
                            resp_result.category,
                            resp_result.confidence,
                        )
            except (AttributeError, ValueError):
                pass

        return response

    async def guarded_generate_content_async(contents, *args, **kwargs):
        text = _extract_text(contents)
        if text.strip():
            result = await guard.scan_async(text, content_type="user_prompt", agent=agent)
            if _should_block(result):
                logger.warning(
                    "burrow blocked vertex async: category=%s confidence=%.2f request_id=%s",
                    result.category,
                    result.confidence,
                    result.request_id,
                )
                raise BurrowBlockedError(result)
            if result.is_warning:
                logger.info(
                    "burrow warn vertex async request: category=%s confidence=%.2f",
                    result.category,
                    result.confidence,
                )

        assert original_generate_async is not None
        response = await original_generate_async(contents, *args, **kwargs)

        if scan_responses:
            try:
                resp_text = response.text
                if resp_text and resp_text.strip():
                    resp_result = await guard.scan_async(
                        resp_text,
                        content_type="tool_response",
                        agent=agent,
                    )
                    if _should_block(resp_result):
                        logger.warning(
                            "burrow blocked vertex async response: category=%s confidence=%.2f",
                            resp_result.category,
                            resp_result.confidence,
                        )
            except (AttributeError, ValueError):
                pass

        return response

    model.generate_content = guarded_generate_content
    if original_generate_async is not None:
        model.generate_content_async = guarded_generate_content_async

    return model
