"""LangChain integration — wrap any LangChain tool with Interceptor guardrails."""

from __future__ import annotations

from typing import Any, Optional, Type, Union

from langchain_core.callbacks import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)
from langchain_core.tools import BaseTool
from pydantic import BaseModel

from interceptor.core import Interceptor
from interceptor.types import DecisionVerdict


class InterceptedTool(BaseTool):
    """A LangChain tool wrapper that runs through an Interceptor guard.

    Usage::

        from interceptor.integrations.langchain import InterceptedTool
        from my_tools import MyTool

        guard = Interceptor(mode="strict")
        safe_tool = InterceptedTool(tool=MyTool(), guard=guard)
    """

    tool: BaseTool
    guard: Interceptor
    user_role: str = "user"

    def __init__(self, tool: BaseTool, guard: Interceptor, user_role: str = "user", **kwargs: Any):
        super().__init__(
            name=tool.name,
            description=tool.description,
            args_schema=tool.args_schema,
            tool=tool,
            guard=guard,
            user_role=user_role,
            **kwargs,
        )

    def _run(
        self,
        *args: Any,
        run_manager: Optional[CallbackManagerForToolRun] = None,
        **kwargs: Any,
    ) -> Any:
        # Evaluate the call
        decision = self.guard.run(self.name, kwargs, user_role=self.user_role)

        if not decision.allowed:
            if decision.decision == DecisionVerdict.BLOCKED:
                return f"Error: This action was BLOCKED by security policy. Reason: {decision.explanation or decision.reason}"
            if decision.decision == DecisionVerdict.CONFIRMATION_REQUIRED:
                return f"Error: This action REQUIRES HUMAN CONFIRMATION. Reason: {decision.explanation or decision.reason}"
            return f"Error: Action not allowed. {decision.reason}"

        # Run the actual tool. We use .run() to let LangChain handle 
        # validation and callbacks normally.
        return self.tool.run(kwargs, run_manager=run_manager, **kwargs)

    async def _arun(
        self,
        *args: Any,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
        **kwargs: Any,
    ) -> Any:
        # Evaluate the call asynchronously
        decision = await self.guard.arun(self.name, kwargs, user_role=self.user_role)

        if not decision.allowed:
            if decision.decision == DecisionVerdict.BLOCKED:
                return f"Error: This action was BLOCKED by security policy. Reason: {decision.explanation or decision.reason}"
            if decision.decision == DecisionVerdict.CONFIRMATION_REQUIRED:
                return f"Error: This action REQUIRES HUMAN CONFIRMATION. Reason: {decision.explanation or decision.reason}"
            return f"Error: Action not allowed. {decision.reason}"

        # Run the actual tool asynchronously
        return await self.tool.arun(kwargs, run_manager=run_manager, **kwargs)
