"""
eaigins_sdk.integrations
────────────────────────
Ready-made governance hooks for the three major agentic frameworks.

    LangChain  → EaignsLangChainCallback  (BaseCallbackHandler subclass)
    CrewAI     → governed_crewai_tool     (wraps @crewai_tool functions)
    AutoGen    → EaignsAutoGenProxy       (UserProxyAgent subclass)

All integrations call ``ControlPlane.evaluate()`` at the tool execution
boundary of their respective framework, before the tool body runs.
"""

from __future__ import annotations
from typing import Any, Callable, Optional

from ..models        import AgentAction, GovernanceViolationError, ExecutionDecision
from ..control_plane import ControlPlane


# ═══════════════════════════════════════════════════════════════════════════════
# LANGCHAIN INTEGRATION
# ═══════════════════════════════════════════════════════════════════════════════

class EaignsLangChainCallback:
    """
    LangChain callback handler that intercepts every tool invocation.

    Plugs into LangChain's ``AgentExecutor`` via the ``callbacks`` parameter.

    Usage::

        from langchain.agents import AgentExecutor, create_react_agent
        from eaigins_sdk.integrations import EaignsLangChainCallback
        from eaigins_sdk import ControlPlane
        from eaigins_sdk.policies import HEALTHCARE_POLICIES

        cp       = ControlPlane(policies=HEALTHCARE_POLICIES)
        callback = EaignsLangChainCallback(cp, agent_id="langchain::my-agent")

        executor = AgentExecutor(
            agent=agent, tools=tools, verbose=True,
            callbacks=[callback],
            handle_parsing_errors=True,
        )

    When the LLM chooses to call a tool, the callback fires BEFORE the
    tool body executes and raises ``GovernanceViolationError`` on BLOCK.
    The LangChain executor surfaces this as a tool error; the agent sees
    "Action blocked by governance" and adjusts its reasoning.
    """

    def __init__(
        self,
        control_plane:   ControlPlane,
        agent_id:        str = "langchain::react-agent",
        context_sources: Optional[dict[str, list[str]]] = None,
    ):
        """
        Args:
            control_plane:   Configured ControlPlane instance.
            agent_id:        Agent identifier stamped on AgentActions.
            context_sources: Optional mapping of tool_name → [source, ...].
                             Used to populate AgentAction.context_sources.
        """
        self.cp              = control_plane
        self.agent_id        = agent_id
        self.context_sources = context_sources or {}
        self._last_reasoning = ""

    # ── LangChain callback hooks ──────────────────────────────────────────────

    def on_agent_action(self, action: Any, **kwargs):
        """Capture the LLM's chain-of-thought reasoning for provenance."""
        self._last_reasoning = str(getattr(action, "log", ""))[:300]

    def on_tool_start(
        self,
        serialized: dict,
        input_str:  str,
        **kwargs,
    ):
        """Intercept the tool call BEFORE execution."""
        tool_name = serialized.get("name", "unknown_tool")
        reasoning = self._last_reasoning or (
            f"LangChain ReAct agent invoking {tool_name}: {str(input_str)[:150]}"
        )

        # Parse tool_params from input_str (may be JSON or plain string)
        import json as _json
        try:
            params = _json.loads(input_str) if isinstance(input_str, str) else {}
        except Exception:
            params = {"input": str(input_str)[:200]}

        action = AgentAction(
            agent_id        = self.agent_id,
            prompt          = f"LangChain agent invoking {tool_name}",
            tool_name       = tool_name,
            tool_params     = params,
            context_sources = self.context_sources.get(tool_name, []),
            reasoning       = reasoning,
            metadata        = {"framework": "langchain"},
        )
        decision = self.cp.evaluate(action)

        if decision.decision == ExecutionDecision.BLOCK:
            raise GovernanceViolationError(decision)
        # ESCALATE: allow through but log (human approval tracked in provenance)

    def on_tool_end(self, output: str, **kwargs):
        """Clear cached reasoning after tool completes."""
        self._last_reasoning = ""

    def on_tool_error(self, error: Exception, **kwargs):
        self._last_reasoning = ""


def _try_register_langchain_base():
    """Attempt to make EaignsLangChainCallback a proper BaseCallbackHandler."""
    try:
        from langchain.callbacks.base import BaseCallbackHandler
        if BaseCallbackHandler not in EaignsLangChainCallback.__bases__:
            EaignsLangChainCallback.__bases__ = (BaseCallbackHandler,)
    except ImportError:
        pass  # langchain not installed — duck-typed callback still works


_try_register_langchain_base()


# ═══════════════════════════════════════════════════════════════════════════════
# CREWAI INTEGRATION
# ═══════════════════════════════════════════════════════════════════════════════

def governed_crewai_tool(
    control_plane:   ControlPlane,
    tool_name:       str,
    context_sources: Optional[list[str]] = None,
    agent_id:        str                 = "crewai::agent",
    metadata:        Optional[dict]      = None,
) -> Callable:
    """
    Decorator for CrewAI ``@tool`` functions that adds Eaigins governance.

    Usage::

        from crewai.tools import tool as crewai_tool
        from eaigins_sdk.integrations import governed_crewai_tool
        from eaigins_sdk import ControlPlane
        from eaigins_sdk.policies import AML_POLICIES

        cp = ControlPlane(policies=AML_POLICIES)

        @crewai_tool("File SAR with Regulator")
        @governed_crewai_tool(cp, tool_name="file_sar",
                              context_sources=["compliance_reports", "kyc_documents"])
        def file_sar(case_id: str, customer_id: str) -> str:
            \"\"\"Submit a SAR to FinCEN. Requires human authorization.\"\"\"
            return submit_to_fincen(case_id, customer_id)

    The Eaigins interception fires BEFORE the CrewAI @tool body.
    If blocked, ``GovernanceViolationError`` propagates as a tool failure.
    The crew's verbose output will show the governance rejection.
    """
    import functools

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Build a prompt from positional + keyword args
            arg_summary = ", ".join(
                list(str(a)[:40] for a in args) +
                [f"{k}={str(v)[:40]}" for k, v in kwargs.items()]
            )
            action = AgentAction(
                agent_id        = agent_id,
                prompt          = f"CrewAI agent invoking {tool_name}({arg_summary})",
                tool_name       = tool_name,
                tool_params     = kwargs,
                context_sources = context_sources or [],
                reasoning       = f"CrewAI task execution: {tool_name}",
                metadata        = metadata or {"framework": "crewai"},
            )
            decision = control_plane.evaluate(action)
            if decision.decision == ExecutionDecision.BLOCK:
                raise GovernanceViolationError(decision)
            return func(*args, **kwargs)

        wrapper._eaigins_cp        = control_plane
        wrapper._eaigins_tool_name = tool_name
        return wrapper

    return decorator


# ═══════════════════════════════════════════════════════════════════════════════
# AUTOGEN INTEGRATION
# ═══════════════════════════════════════════════════════════════════════════════

class EaignsAutoGenProxy:
    """
    AutoGen UserProxyAgent subclass that intercepts function calls via Eaigins.

    Usage::

        import autogen
        from eaigins_sdk.integrations import EaignsAutoGenProxy
        from eaigins_sdk import ControlPlane
        from eaigins_sdk.policies import CUSTOMERCARE_POLICIES

        cp = ControlPlane(policies=CUSTOMERCARE_POLICIES)

        def process_refund(customer_id: str, amount: float, reason: str) -> str:
            return real_refund_api(customer_id, amount, reason)

        proxy = EaignsAutoGenProxy.create(
            cp,
            function_map={
                "process_refund":   process_refund,
                "export_data":      export_data,
            },
            tool_context_map={
                "process_refund":   ["payment_records", "order_management_db"],
                "export_data":      ["crm_customer_profile"],
            },
            agent_id = "autogen::cx-agent",
        )

    Every function in ``function_map`` is wrapped: Eaigins evaluates the
    action before the function body runs.  BLOCK raises GovernanceViolationError
    which AutoGen surfaces as a function execution error to the assistant.
    """

    @staticmethod
    def create(
        control_plane:    ControlPlane,
        function_map:     dict[str, Callable],
        tool_context_map: Optional[dict[str, list[str]]] = None,
        agent_id:         str = "autogen::agent",
        name:             str = "eaigins_proxy",
        **proxy_kwargs,
    ):
        """
        Build a governed AutoGen UserProxyAgent.

        Returns a ``UserProxyAgent`` (or duck-typed fallback) whose
        ``function_map`` is wrapped with Eaigins governance.
        """
        try:
            import autogen
        except ImportError:
            raise ImportError(
                "pyautogen is required for AutoGen integration: pip install pyautogen"
            )

        context_map = tool_context_map or {}
        last_msg: dict = {"content": ""}

        governed_map: dict[str, Callable] = {}
        for fn_name, fn in function_map.items():
            def make_wrapper(name: str, func: Callable) -> Callable:
                import functools
                @functools.wraps(func)
                def wrapper(**kwargs):
                    action = AgentAction(
                        agent_id        = agent_id,
                        prompt          = last_msg.get("content", "")[:200] or
                                          f"AutoGen agent invoking {name}",
                        tool_name       = name,
                        tool_params     = kwargs,
                        context_sources = context_map.get(name, []),
                        reasoning       = last_msg.get("content", "")[:200],
                        metadata        = {"framework": "autogen"},
                    )
                    decision = control_plane.evaluate(action)
                    if decision.decision == ExecutionDecision.BLOCK:
                        raise GovernanceViolationError(decision)
                    return func(**kwargs)
                return wrapper
            governed_map[fn_name] = make_wrapper(fn_name, fn)

        class _EaignsProxy(autogen.UserProxyAgent):
            def execute_function(self, func_call):
                # Capture assistant reasoning before execution
                msgs = self.chat_messages
                for agent_msgs in msgs.values():
                    if agent_msgs:
                        last_msg["content"] = (
                            agent_msgs[-1].get("content", "") or ""
                        )[:300]
                return super().execute_function(func_call)

        return _EaignsProxy(
            name=name,
            function_map=governed_map,
            human_input_mode=proxy_kwargs.pop("human_input_mode", "NEVER"),
            code_execution_config=proxy_kwargs.pop("code_execution_config", False),
            **proxy_kwargs,
        )
