"""
CrewAI Integration for Agendex Governance.

Provides GovernedCrew - a wrapper for CrewAI's Crew that automatically
routes all tool calls through Agendex.

Example:
    from crewai import Agent, Task, Crew
    from agendex import AgendexClient
    from agendex.integrations.crewai import GovernedCrew
    
    # Create your agents and tasks as normal
    researcher = Agent(role="Researcher", tools=[search_tool])
    writer = Agent(role="Writer", tools=[write_tool])
    
    task = Task(description="Research and write about AI")
    
    # Replace Crew with GovernedCrew
    client = AgendexClient()
    crew = GovernedCrew(
        agents=[researcher, writer],
        tasks=[task],
        agendex=client,
        task="content_creation"  # Agendex task context
    )
    
    # Run as normal - all tool calls are governed
    result = crew.kickoff()
"""
from __future__ import annotations

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

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

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

# Lazy import CrewAI to avoid hard dependency
try:
    from crewai import Agent, Crew, Task
    from crewai.tools import BaseTool as CrewAIBaseTool
    from pydantic import BaseModel, Field
    CREWAI_AVAILABLE = True
except ImportError:
    CREWAI_AVAILABLE = False
    Crew = object  # type: ignore
    Agent = object  # type: ignore
    Task = object  # type: ignore


def _create_governed_crewai_tool(
    tool: Any,
    agendex: AgendexClient,
    task: str,
    agent_role: str,
    on_event: Optional[Callable[[Dict], None]] = None,
    reasoning_state: Optional[Dict[str, Any]] = None,
) -> Any:
    """
    Create a new CrewAI tool that wraps the original with governance.
    """
    if not CREWAI_AVAILABLE:
        return tool
    
    tool_name = getattr(tool, "name", "unknown_tool")
    tool_description = getattr(tool, "description", "")
    args_schema = getattr(tool, "args_schema", None)
    
    # Store reference to original tool for execution
    original_tool = tool
    
    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
        
        # Include agent role and latest reasoning in context
        context = {"agent_role": agent_role}
        reasoning = reasoning_state.get("current_reasoning") if reasoning_state else None
        if reasoning:
            context["reasoning"] = reasoning
        
        try:
            result = agendex.invoke(
                action=action,
                params=params,
                task=task,
                context=context,
                on_event=on_event,
            )
            
            # Shadow mode: execute locally using the original tool
            if isinstance(result, dict) and result.get("_agendex") == "shadow":
                # Call the original tool's _run method with proper binding
                if hasattr(original_tool, '_run'):
                    return original_tool._run(*args, **kwargs)
                elif hasattr(original_tool, 'func') and original_tool.func:
                    return original_tool.func(*args, **kwargs)
                elif callable(original_tool):
                    return original_tool(*args, **kwargs)
                else:
                    return f"Error: Could not execute original tool {tool_name}"
            
            # Enforce mode: result from /invoke
            return result
            
        except DeniedError as e:
            logger.warning(f"Tool {tool_name} denied for agent {agent_role}: {e.reason}")
            return f"[GOVERNANCE DENIED] Action '{action}' was blocked: {e.reason}"
        
        except PendingApprovalError as e:
            logger.info(f"Tool {tool_name} requires approval for agent {agent_role}")
            return f"[APPROVAL REQUIRED] Action '{action}' requires human approval. Approval ID: {e.approval_id}"
        finally:
            if reasoning_state is not None:
                reasoning_state["current_reasoning"] = None
    
    # Create a new tool class dynamically
    class GovernedTool(CrewAIBaseTool):
        name: str = tool_name
        description: str = tool_description
        
        def _run(self, *args: Any, **kwargs: Any) -> Any:
            return governed_func(*args, **kwargs)
    
    # If original has args_schema, try to use it
    if args_schema:
        GovernedTool.args_schema = args_schema
    
    return GovernedTool()


def _wrap_agent_tools(
    agent: Agent,
    agendex: AgendexClient,
    task: str,
    on_event: Optional[Callable[[Dict], None]] = None,
    reasoning_state: Optional[Dict[str, Any]] = None,
) -> Agent:
    """
    Create governed versions of all tools in a CrewAI Agent.
    """
    agent_role = agent.role
    
    if not hasattr(agent, "tools") or not agent.tools:
        return agent
    
    wrapped_tools = []
    for tool in agent.tools:
        governed_tool = _create_governed_crewai_tool(
            tool=tool,
            agendex=agendex,
            task=task,
            agent_role=agent_role,
            on_event=on_event,
            reasoning_state=reasoning_state,
        )
        wrapped_tools.append(governed_tool)
    
    agent.tools = wrapped_tools
    logger.debug(f"Wrapped {len(wrapped_tools)} tools for agent '{agent_role}'")
    
    return agent


class GovernedCrew:
    """
    Wrapper for CrewAI's Crew with Agendex governance.
    
    Automatically:
    1. Wraps all agent tools to route through Agendex
    2. Captures agent context (role, goal) for governance
    3. Handles governance decisions (deny, approval required)
    
    Usage:
        crew = GovernedCrew(
            agents=[agent1, agent2],
            tasks=[task],
            agendex=AgendexClient(),
            task="my_workflow"
        )
        result = crew.kickoff()
    """
    
    def __init__(
        self,
        agents: List[Agent],
        tasks: List[Task],
        agendex: AgendexClient,
        task: str,
        on_event: Optional[Callable[[Dict], None]] = None,
        **kwargs: Any,
    ):
        """
        Initialize governed crew.
        
        Args:
            agents: List of CrewAI agents
            tasks: List of CrewAI tasks
            agendex: AgendexClient instance
            task: Agendex task name for governance context
            on_event: Optional callback for governance events
            **kwargs: Additional arguments passed to Crew
        """
        if not CREWAI_AVAILABLE:
            raise ImportError(
                "CrewAI is required for GovernedCrew. "
                "Install with: pip install agendex[crewai]"
            )
        
        self._agendex = agendex
        self._agendex_task = task
        self._on_event = on_event
        
        # Wrap all agent tools for governance, attach reasoning capture if available
        for agent in agents:
            reasoning_state: Dict[str, Any] = {"current_reasoning": None}
            original_step_callback = getattr(agent, "step_callback", None)

            def _step_callback(step_output, *, _state=reasoning_state, _orig=original_step_callback, _role=agent.role):
                try:
                    thought = getattr(step_output, "thought", None)
                    if thought:
                        _state["current_reasoning"] = thought
                        if on_event:
                            on_event({
                                "type": "reasoning",
                                "agent": _role,
                                "content": thought,
                            })
                except Exception:
                    pass
                if _orig:
                    try:
                        _orig(step_output)
                    except Exception:
                        pass

            if hasattr(agent, "step_callback"):
                agent.step_callback = _step_callback

            _wrap_agent_tools(
                agent=agent,
                agendex=agendex,
                task=task,
                on_event=on_event,
                reasoning_state=reasoning_state,
            )
        
        # Create the underlying Crew
        self._crew = Crew(
            agents=agents,
            tasks=tasks,
            **kwargs,
        )
        
        logger.info(
            f"GovernedCrew initialized: task={task}, "
            f"agents={[a.role for a in agents]}, mode={agendex.mode}"
        )
    
    def kickoff(self, inputs: Optional[Dict[str, Any]] = None) -> Any:
        """
        Start the crew with governance.
        """
        if self._on_event:
            self._on_event({
                "type": "crew_start",
                "task": self._agendex_task,
                "agents": [a.role for a in self._crew.agents],
            })
        
        try:
            result = self._crew.kickoff(inputs=inputs)
            
            if self._on_event:
                self._on_event({
                    "type": "crew_complete",
                    "task": self._agendex_task,
                    "success": True,
                })
            
            return result
            
        except Exception as e:
            if self._on_event:
                self._on_event({
                    "type": "crew_complete",
                    "task": self._agendex_task,
                    "success": False,
                    "error": str(e),
                })
            raise
    
    @property
    def agents(self) -> List[Agent]:
        """Get the crew's agents."""
        return self._crew.agents
    
    @property
    def tasks(self) -> List[Task]:
        """Get the crew's tasks."""
        return self._crew.tasks
    
    def __getattr__(self, name: str) -> Any:
        """Delegate attribute access to the underlying crew."""
        return getattr(self._crew, name)


# Convenience function for quick setup
def govern_crew(
    crew: Crew,
    agendex: AgendexClient,
    task: str,
    on_event: Optional[Callable[[Dict], None]] = None,
) -> GovernedCrew:
    """
    Convert an existing Crew to a GovernedCrew.
    
    Args:
        crew: Existing CrewAI Crew
        agendex: AgendexClient instance
        task: Agendex task name for governance context
        on_event: Optional callback for governance events
    
    Returns:
        GovernedCrew wrapping the same agents and tasks
    """
    if not CREWAI_AVAILABLE:
        raise ImportError(
            "CrewAI is required for govern_crew. "
            "Install with: pip install agendex[crewai]"
        )
    
    return GovernedCrew(
        agents=crew.agents,
        tasks=crew.tasks,
        agendex=agendex,
        task=task,
        on_event=on_event,
        process=crew.process,
        verbose=crew.verbose,
    )
