"""Guardrail models, exceptions, and runner helpers for the Swarm framework.

Guardrails validate input before the first LLM call and output after the loop
exits.  Unlike lifecycle hooks, guardrails **can** halt execution via a tripwire.
"""

import logging
from typing import Any, Awaitable, Callable

from pydantic import BaseModel

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------


class GuardrailResult(BaseModel):
    """Result from a guardrail function."""

    guardrail_name: str = ""
    tripwire_triggered: bool = False
    output_info: Any = None


class InputGuardrail(BaseModel):
    """Validates input before the first LLM call."""

    name: str
    guardrail_fn: Callable[..., Awaitable[GuardrailResult]]

    class Config:
        arbitrary_types_allowed = True


class OutputGuardrail(BaseModel):
    """Validates final output after the loop terminates."""

    name: str
    guardrail_fn: Callable[..., Awaitable[GuardrailResult]]

    class Config:
        arbitrary_types_allowed = True


class OutputGuardrailContext(BaseModel):
    """Data passed to output guardrail functions."""

    agent_name: str
    agent_output: str
    context_variables: dict[str, Any]


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------


class Tripwire(Exception):
    """Fast-abort sugar for guardrails.

    Moved from ``swarm.hooks`` — tripwires belong to the guardrail concern.
    """

    def __init__(self, action: str, payload: Any = None) -> None:
        self.action = action
        self.payload = payload
        super().__init__(f"Tripwire: {action}")


class InputGuardrailTripped(Exception):
    """Raised when an input guardrail trips."""

    def __init__(self, guardrail_name: str, result: GuardrailResult) -> None:
        self.guardrail_name = guardrail_name
        self.result = result
        super().__init__(f"Input guardrail tripped: {guardrail_name}")


class OutputGuardrailTripped(Exception):
    """Raised when an output guardrail trips (batch mode only)."""

    def __init__(self, guardrail_name: str, result: GuardrailResult) -> None:
        self.guardrail_name = guardrail_name
        self.result = result
        super().__init__(f"Output guardrail tripped: {guardrail_name}")


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


def _extract_last_assistant_text(
    message_history: list[Any],
) -> str:
    """Scan *message_history* backward for the last assistant message and
    extract concatenated text content.

    Algorithm:
    1. Walk backward from the end of ``message_history``.
    2. Find the first entry with ``role == "assistant"``.
    3. Extract ``content`` (expected shape: ``list[dict]`` per ``message.py:21``).
    4. Select only ``{"type": "text", "text": ...}`` parts.
    5. Concatenate their ``"text"`` values with ``"\\n"``.
    6. Default to ``""`` if nothing found.
    """
    for entry in reversed(message_history):
        role = (
            entry.get("role")
            if isinstance(entry, dict)
            else getattr(entry, "role", None)
        )
        if role in ("assistant", "ASSISTANT"):
            content = (
                entry.get("content")
                if isinstance(entry, dict)
                else getattr(entry, "content", None)
            )
            if not isinstance(content, list):
                return str(content) if content else ""
            text_parts: list[str] = []
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    text_parts.append(part.get("text", ""))
            return "\n".join(text_parts)
    return ""


async def run_input_guardrails(
    agent: Any,
    messages: list[Any],
    context_variables: dict[str, Any],
) -> list[GuardrailResult]:
    """Run input guardrails sequentially.  First trip raises ``InputGuardrailTripped``."""
    guardrails: list[Any] = getattr(agent, "input_guardrails", None) or []
    results: list[GuardrailResult] = []
    for guardrail in guardrails:
        if not isinstance(guardrail, InputGuardrail):
            logger.warning(
                "Skipping invalid input guardrail: %s", type(guardrail).__name__
            )
            continue
        result = await guardrail.guardrail_fn(agent, messages, context_variables)
        result.guardrail_name = guardrail.name
        results.append(result)
        if result.tripwire_triggered:
            raise InputGuardrailTripped(guardrail.name, result)
    return results


async def run_output_guardrails(
    agent: Any,
    output_context: OutputGuardrailContext,
) -> list[GuardrailResult]:
    """Run output guardrails sequentially.  First trip halts remaining guardrails
    and returns the results collected so far (caller decides whether to raise).
    """
    guardrails: list[Any] = getattr(agent, "output_guardrails", None) or []
    results: list[GuardrailResult] = []
    for guardrail in guardrails:
        if not isinstance(guardrail, OutputGuardrail):
            logger.warning(
                "Skipping invalid output guardrail: %s", type(guardrail).__name__
            )
            continue
        result = await guardrail.guardrail_fn(agent, output_context)
        result.guardrail_name = guardrail.name
        results.append(result)
        if result.tripwire_triggered:
            break
    return results
