"""AutoPIL LangChainGuard. See packages/core/autopil/langchain_guard.py for full docs."""

from __future__ import annotations
from typing import Any, Optional
from .guard import ContextGuard
from .models import SensitivityLevel


class LangChainGuard(ContextGuard):
    """ContextGuard for LangChain agents. Tags audit events with source_type='langchain'."""

    _source_type: str = "langchain"

    def wrap_tool(
        self,
        tool: Any,
        *,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ) -> Any:
        """Wrap a LangChain BaseTool._run with AutoPIL policy enforcement."""
        original_run = tool._run
        tool._run = self.protect(
            agent_role=agent_role,
            user_id=user_id,
            source_id=source_id,
            sensitivity_level=sensitivity_level,
            session_id=session_id,
        )(original_run)
        return tool

    def wrap_tool_async(
        self,
        tool: Any,
        *,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ) -> Any:
        """Wrap a LangChain BaseTool._arun with AutoPIL async policy enforcement."""
        if hasattr(tool, "_arun") and callable(tool._arun):
            original_arun = tool._arun
        else:
            _sync_run = tool._run
            async def original_arun(*args: Any, **kwargs: Any) -> Any:
                return _sync_run(*args, **kwargs)
        tool._arun = self.protect_async(
            agent_role=agent_role,
            user_id=user_id,
            source_id=source_id,
            sensitivity_level=sensitivity_level,
            session_id=session_id,
        )(original_arun)
        return tool
