"""AutoGen integration — function wrapping and message checking."""

from __future__ import annotations

import functools
from typing import TYPE_CHECKING, Any, Callable

if TYPE_CHECKING:
    from ..client import Guard


class SurfinguardAutoGenGuard:
    """AutoGen integration for Surfinguard security checks.

    Usage:
        from surfinguard import Guard
        from surfinguard.integrations.autogen import SurfinguardAutoGenGuard

        guard = Guard(api_key="sg_live_...")
        ag_guard = SurfinguardAutoGenGuard(guard)

        @ag_guard.wrap_function("command")
        def execute_command(command: str) -> str:
            ...

        # Or check messages
        ag_guard.check_message(message)
    """

    def __init__(self, guard: Guard) -> None:
        self._guard = guard

    def wrap_function(self, action_type: str = "command") -> Callable[..., Any]:
        """Decorator that checks the first argument before function execution.

        Args:
            action_type: The action type to check.

        Returns:
            A decorator function.
        """
        guard = self._guard

        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
            @functools.wraps(fn)
            def wrapper(*args: Any, **kwargs: Any) -> Any:
                value: str | None = None
                if args:
                    value = str(args[0])
                elif kwargs:
                    first_key = next(iter(kwargs))
                    value = str(kwargs[first_key])

                if value is not None:
                    guard.check(action_type, value)

                return fn(*args, **kwargs)

            return wrapper

        return decorator

    def check_message(self, message: Any) -> None:
        """Check a message for prompt injection or manipulation.

        Handles both string messages and dict messages with a 'content' key.

        Args:
            message: A string or dict with 'content' key.
        """
        content: str | None = None

        if isinstance(message, str):
            content = message
        elif isinstance(message, dict):
            content = message.get("content")
            if content and not isinstance(content, str):
                content = str(content)

        if content:
            self._guard.check("text", content)
