"""CrewAI integration — tool decorator and step callback."""

from __future__ import annotations

import functools
from typing import TYPE_CHECKING, Any, Callable

if TYPE_CHECKING:
    from ..client import Guard


class SurfinguardCrewGuard:
    """CrewAI integration for Surfinguard security checks.

    Usage:
        from surfinguard import Guard
        from surfinguard.integrations.crewai import SurfinguardCrewGuard

        guard = Guard(api_key="sg_live_...")
        crew_guard = SurfinguardCrewGuard(guard)

        @crew_guard.check_tool("command")
        def execute_command(command: str) -> str:
            ...

        # Or use as a step callback
        crew = Crew(..., step_callback=crew_guard.step_callback)
    """

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

    def check_tool(self, action_type: str = "command") -> Callable[..., Any]:
        """Decorator that checks the first argument before tool 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 step_callback(self, step_output: Any) -> None:
        """CrewAI step callback that checks tool inputs.

        Attach to a Crew: `Crew(..., step_callback=crew_guard.step_callback)`

        Args:
            step_output: The step output object from CrewAI.
        """
        # CrewAI step_output has tool_input when a tool is invoked
        tool_input = getattr(step_output, "tool_input", None)
        if tool_input and isinstance(tool_input, str):
            self._guard.check("command", tool_input)
