"""
LangChain Integration for Agendex Governance.

Provides GovernedAgentExecutor - a wrapper for LangChain's AgentExecutor
that automatically routes all tool calls through Agendex.

Example:
    from langchain.agents import create_react_agent, AgentExecutor
    from langchain_openai import ChatOpenAI
    from agendex import AgendexClient
    from agendex.integrations.langchain import GovernedAgentExecutor
    
    # Create your agent as normal
    llm = ChatOpenAI()
    agent = create_react_agent(llm, tools, prompt)
    
    # Wrap with GovernedAgentExecutor
    client = AgendexClient()
    executor = GovernedAgentExecutor(
        agent=agent,
        tools=tools,
        agendex=client,
        task="my_task"
    )
    
    # Run as normal - all tool calls are now governed
    result = executor.invoke({"input": "What's the revenue for Q1?"})
"""
from __future__ import annotations

import functools
import logging
from typing import Any, Callable, Dict, List, Optional, Sequence, Type

from ..client import AgendexClient
from ..errors import DeniedError, PendingApprovalError

logger = logging.getLogger("agendex.langchain")

# Lazy import LangChain to avoid hard dependency
try:
    from langchain.agents import AgentExecutor
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.agents import AgentAction
    from langchain_core.tools import BaseTool, StructuredTool
    from pydantic import BaseModel
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    AgentExecutor = object  # type: ignore
    BaseTool = object  # type: ignore


class _ReasoningCapture(BaseCallbackHandler if LANGCHAIN_AVAILABLE else object):
    """
    LangChain callback handler that captures agent reasoning.
    """
    
    def __init__(self, state: Dict[str, Any]):
        super().__init__()
        self._state = state
    
    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: Any = None,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        """Capture reasoning from agent action log."""
        if hasattr(action, "log") and action.log:
            self._state["current_reasoning"] = action.log
            on_event = self._state.get("on_event")
            if on_event:
                on_event({
                    "type": "reasoning",
                    "content": action.log,
                })
    
    def on_tool_end(
        self,
        output: str,
        *,
        run_id: Any = None,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        """Tool finished - clear current thought."""
        self._state["current_reasoning"] = None


def _create_governed_tool(
    tool: BaseTool,
    agendex: AgendexClient,
    task: str,
    state: Dict[str, Any],
) -> BaseTool:
    """
    Create a new tool that wraps the original with governance.
    
    Uses StructuredTool.from_function to create a proper LangChain tool
    that routes through Agendex.
    """
    tool_name = tool.name
    tool_description = tool.description
    original_func = tool.func if hasattr(tool, 'func') else None
    args_schema = tool.args_schema if hasattr(tool, 'args_schema') else None
    
    def governed_func(*args: Any, **kwargs: Any) -> Any:
        """Governed wrapper that routes through Agendex."""
        action = f"tool.{tool_name}"
        
        # Build params
        params = {"kwargs": kwargs}
        if args:
            params["args"] = args
        
        # Get current reasoning context
        reasoning = state.get("current_reasoning")
        context = {"reasoning": reasoning} if reasoning else {}
        
        on_event = state.get("on_event")
        
        try:
            result = agendex.invoke(
                action=action,
                params=params,
                task=task,
                context=context,
                on_event=on_event,
            )
            
            # Shadow mode: execute locally
            if isinstance(result, dict) and result.get("_agendex") == "shadow":
                # Execute the original tool
                if original_func:
                    return original_func(*args, **kwargs)
                else:
                    # Fall back to tool.invoke for tools without func
                    return tool.invoke(kwargs if kwargs else (args[0] if args else ""))
            
            # Enforce mode: result from /invoke
            return result
            
        except DeniedError as e:
            logger.warning(f"Tool {tool_name} denied: {e.reason}")
            return f"[GOVERNANCE DENIED] Action '{action}' was blocked: {e.reason}"
        
        except PendingApprovalError as e:
            logger.info(f"Tool {tool_name} requires approval: {e.approval_id}")
            return f"[APPROVAL REQUIRED] Action '{action}' requires human approval. Approval ID: {e.approval_id}"
    
    # Create a new StructuredTool with the governed function
    governed_tool = StructuredTool.from_function(
        func=governed_func,
        name=tool_name,
        description=tool_description,
        args_schema=args_schema,
        return_direct=getattr(tool, 'return_direct', False),
    )
    
    return governed_tool


class GovernedAgentExecutor:
    """
    Wrapper for LangChain's AgentExecutor with Agendex governance.
    
    Automatically:
    1. Wraps all tools to route through Agendex
    2. Captures agent reasoning for audit context
    3. Handles governance decisions (deny, approval required)
    
    Usage:
        executor = GovernedAgentExecutor(
            agent=agent,
            tools=tools,
            agendex=AgendexClient(),
            task="my_task"
        )
        result = executor.invoke({"input": "..."})
    """
    
    def __init__(
        self,
        agent: Any,
        tools: Sequence[BaseTool],
        agendex: AgendexClient,
        task: str,
        on_event: Optional[Callable[[Dict], None]] = None,
        **kwargs: Any,
    ):
        """
        Initialize governed agent executor.
        
        Args:
            agent: The LangChain agent (from create_react_agent, etc.)
            tools: List of LangChain tools
            agendex: AgendexClient instance
            task: Task name for governance context
            on_event: Optional callback for governance events
            **kwargs: Additional arguments passed to AgentExecutor
        """
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is required for GovernedAgentExecutor. "
                "Install with: pip install agendex[langchain]"
            )
        
        self._agendex = agendex
        self._task = task
        
        # Shared state for reasoning capture
        self._state: Dict[str, Any] = {
            "current_reasoning": None,
            "on_event": on_event,
        }
        
        # Create reasoning capture callback
        reasoning_callback = _ReasoningCapture(self._state)
        
        # Add our callback to existing callbacks
        existing_callbacks = list(kwargs.pop("callbacks", []) or [])
        existing_callbacks.append(reasoning_callback)
        
        # Create governed versions of all tools
        wrapped_tools = []
        for tool in tools:
            governed_tool = _create_governed_tool(
                tool=tool,
                agendex=agendex,
                task=task,
                state=self._state,
            )
            wrapped_tools.append(governed_tool)
        
        # Create the underlying AgentExecutor with wrapped tools
        self._executor = AgentExecutor(
            agent=agent,
            tools=wrapped_tools,
            callbacks=existing_callbacks,
            **kwargs,
        )
        
        logger.info(
            f"GovernedAgentExecutor initialized: task={task}, "
            f"tools={[t.name for t in tools]}, mode={agendex.mode}"
        )
    
    def invoke(self, input: Dict[str, Any], **kwargs: Any) -> Dict[str, Any]:
        """Run the agent with governance."""
        return self._executor.invoke(input, **kwargs)
    
    async def ainvoke(self, input: Dict[str, Any], **kwargs: Any) -> Dict[str, Any]:
        """Run the agent asynchronously with governance."""
        return await self._executor.ainvoke(input, **kwargs)
    
    @property
    def agent(self) -> Any:
        """Get the underlying agent."""
        return self._executor.agent
    
    @property
    def tools(self) -> List[BaseTool]:
        """Get the wrapped tools."""
        return self._executor.tools
    
    @property
    def verbose(self) -> bool:
        """Get verbose setting."""
        return self._executor.verbose
    
    def __getattr__(self, name: str) -> Any:
        """Delegate attribute access to the underlying executor."""
        return getattr(self._executor, name)


# Convenience function for quick setup
def govern_agent(
    executor: AgentExecutor,
    agendex: AgendexClient,
    task: str,
    on_event: Optional[Callable[[Dict], None]] = None,
) -> GovernedAgentExecutor:
    """
    Convert an existing AgentExecutor to a GovernedAgentExecutor.
    
    Args:
        executor: Existing LangChain AgentExecutor
        agendex: AgendexClient instance  
        task: Task name for governance context
        on_event: Optional callback for governance events
    
    Returns:
        GovernedAgentExecutor wrapping the same agent and tools
    """
    if not LANGCHAIN_AVAILABLE:
        raise ImportError(
            "LangChain is required for govern_agent. "
            "Install with: pip install agendex[langchain]"
        )
    
    return GovernedAgentExecutor(
        agent=executor.agent,
        tools=executor.tools,
        agendex=agendex,
        task=task,
        on_event=on_event,
        verbose=executor.verbose,
        max_iterations=executor.max_iterations,
        max_execution_time=executor.max_execution_time,
        early_stopping_method=executor.early_stopping_method,
        handle_parsing_errors=executor.handle_parsing_errors,
    )
