"""AutoGen integration for Sentrial.

Automatically captures agent conversations, tool calls, and token usage from
Microsoft AutoGen multi-agent systems.

Supports multiple AutoGen versions:
- autogen-agentchat >= 0.4.0 (new event-driven API)
- pyautogen >= 0.2.0 (legacy API)

Usage with AutoGen 0.4+ (autogen-agentchat):

    from autogen_agentchat.agents import AssistantAgent
    from autogen_agentchat.teams import RoundRobinGroupChat
    from sentrial import SentrialClient
    from sentrial.autogen import SentrialAutogenHandler
    
    client = SentrialClient()
    session_id = client.create_session(name="Multi-Agent Chat", agent_name="autogen-team")
    handler = SentrialAutogenHandler(client, session_id)
    
    # Wrap your team with the handler
    team = RoundRobinGroupChat([agent1, agent2])
    result = await handler.run_team(team, task="Analyze this data...")
    
    handler.finish(success=True)

Usage with legacy pyautogen:

    from autogen import AssistantAgent, UserProxyAgent
    from sentrial import SentrialClient
    from sentrial.autogen import SentrialAutogenHandler
    
    client = SentrialClient()
    session_id = client.create_session(name="AutoGen Chat", agent_name="autogen")
    handler = SentrialAutogenHandler(client, session_id)
    
    # Register the handler with agents
    assistant = AssistantAgent("assistant", llm_config=llm_config)
    user_proxy = UserProxyAgent("user_proxy")
    
    handler.register_agents([assistant, user_proxy])
    
    user_proxy.initiate_chat(assistant, message="Hello!")
    
    handler.finish(success=True)
"""

from typing import Any, Dict, List, Optional, Callable, Union
import time
import json

from .client import SentrialClient
from .pricing import MODEL_PRICING, calculate_cost


# Version detection and imports
AUTOGEN_VERSION = "none"

# Try new autogen-agentchat API first (0.4+)
try:
    from autogen_agentchat.base import TaskResult
    from autogen_agentchat.messages import (
        AgentEvent,
        ChatMessage,
        TextMessage,
        ToolCallRequestEvent,
        ToolCallExecutionEvent,
    )
    AUTOGEN_VERSION = "agentchat"
except ImportError:
    pass

# Try legacy pyautogen API
if AUTOGEN_VERSION == "none":
    try:
        from autogen import ConversableAgent
        AUTOGEN_VERSION = "legacy"
    except ImportError:
        pass


class SentrialAutogenHandler:
    """
    AutoGen handler for Sentrial performance monitoring.
    
    Automatically tracks:
    - Agent messages and conversations
    - Tool/function executions
    - Token usage and costs
    - Multi-agent interactions
    
    Works with both AutoGen 0.4+ (autogen-agentchat) and legacy pyautogen.
    """
    
    def __init__(
        self,
        client: SentrialClient,
        session_id: str,
        verbose: bool = False,
    ):
        """
        Initialize Sentrial AutoGen handler.
        
        Args:
            client: SentrialClient instance
            session_id: Active session ID
            verbose: Print tracking info
        """
        self.client = client
        self.session_id = session_id
        self.verbose = verbose
        
        # Tracking state
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.message_count = 0
        self.tool_calls = 0
        
        # Token and cost tracking
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        self.total_tokens = 0
        self.total_cost = 0.0
        self.model_name: Optional[str] = None
        
        # User input/output
        self.user_input: Optional[str] = None
        self.assistant_output: Optional[str] = None
        
        # Agent registry (for legacy API)
        self._registered_agents: List[Any] = []
        self._original_callbacks: Dict[str, Callable] = {}
    
    def _log(self, message: str):
        """Log if verbose mode is enabled."""
        if self.verbose:
            print(f"[Sentrial/AutoGen] {message}")
    
    @property
    def duration_ms(self) -> int:
        """Get duration in milliseconds."""
        if self.start_time is None:
            return 0
        end = self.end_time if self.end_time else time.time()
        return int((end - self.start_time) * 1000)
    
    def set_input(self, user_input: str) -> None:
        """Set the user's original input for this session."""
        self.user_input = user_input
        self._log(f"User input set: {user_input[:50]}...")
    
    # ===== AutoGen 0.4+ (autogen-agentchat) API =====
    
    async def run_team(
        self,
        team: Any,  # BaseGroupChat or similar
        task: str,
        **kwargs
    ) -> Any:
        """
        Run an AutoGen team with Sentrial tracking.
        
        Args:
            team: AutoGen team (RoundRobinGroupChat, SelectorGroupChat, etc.)
            task: The task/query to run
            **kwargs: Additional arguments passed to team.run()
            
        Returns:
            TaskResult from the team execution
        """
        if AUTOGEN_VERSION != "agentchat":
            raise ImportError(
                "run_team() requires autogen-agentchat >= 0.4.0. "
                "Install with: pip install autogen-agentchat>=0.4.0"
            )
        
        self.start_time = time.time()
        self.user_input = task
        
        self._log(f"Starting team run: {task[:50]}...")
        
        # Run with streaming to capture events
        result = None
        async for event in team.run_stream(task=task, **kwargs):
            self._process_agentchat_event(event)
            if isinstance(event, TaskResult):
                result = event
        
        self.end_time = time.time()
        
        if result and result.messages:
            # Capture final output
            last_message = result.messages[-1]
            if hasattr(last_message, 'content'):
                self.assistant_output = str(last_message.content)
        
        return result
    
    def _process_agentchat_event(self, event: Any) -> None:
        """Process an event from autogen-agentchat stream."""
        self._log(f"Event: {type(event).__name__}")
        
        # Handle different event types
        if hasattr(event, '__class__'):
            event_type = event.__class__.__name__
            
            if event_type == "TextMessage" or event_type == "ChatMessage":
                self._track_message(event)
            elif event_type == "ToolCallRequestEvent":
                self._track_tool_request(event)
            elif event_type == "ToolCallExecutionEvent":
                self._track_tool_execution(event)
            elif event_type == "ModelClientResponseEvent":
                self._track_model_response(event)
    
    def _track_message(self, message: Any) -> None:
        """Track a chat message."""
        self.message_count += 1
        
        source = getattr(message, 'source', 'unknown')
        content = getattr(message, 'content', str(message))
        
        self._log(f"Message from {source}: {str(content)[:50]}...")
        
        self.client.track_decision(
            session_id=self.session_id,
            reasoning=f"[{source}] {str(content)[:200]}",
            confidence=1.0
        )
    
    def _track_tool_request(self, event: Any) -> None:
        """Track a tool call request."""
        self.tool_calls += 1
        
        # Extract tool info
        content = getattr(event, 'content', [])
        if isinstance(content, list):
            for call in content:
                tool_name = getattr(call, 'name', 'unknown')
                tool_args = getattr(call, 'arguments', '{}')
                
                self._log(f"Tool request: {tool_name}")
                
                # Parse arguments
                try:
                    args_dict = json.loads(tool_args) if isinstance(tool_args, str) else tool_args
                except json.JSONDecodeError:
                    args_dict = {"raw": tool_args}
                
                self.client.track_tool_call(
                    session_id=self.session_id,
                    tool_name=tool_name,
                    tool_input=args_dict,
                    tool_output={"status": "requested"},
                )
    
    def _track_tool_execution(self, event: Any) -> None:
        """Track a tool execution result."""
        content = getattr(event, 'content', [])
        if isinstance(content, list):
            for result in content:
                tool_name = getattr(result, 'call_id', 'unknown')
                output = getattr(result, 'content', str(result))
                
                self._log(f"Tool result: {tool_name}")
                
                self.client.track_tool_call(
                    session_id=self.session_id,
                    tool_name=f"result_{tool_name}",
                    tool_input={},
                    tool_output={"result": str(output)[:500]},
                )
    
    def _track_model_response(self, event: Any) -> None:
        """Track model response with usage stats."""
        # Try to extract usage from the response
        response = getattr(event, 'response', None)
        if response:
            usage = getattr(response, 'usage', None)
            if usage:
                prompt_tokens = getattr(usage, 'prompt_tokens', 0)
                completion_tokens = getattr(usage, 'completion_tokens', 0)
                
                self.total_prompt_tokens += prompt_tokens
                self.total_completion_tokens += completion_tokens
                self.total_tokens += prompt_tokens + completion_tokens
                
                # Get model name
                model = getattr(response, 'model', None)
                if model:
                    self.model_name = model
                    cost = calculate_cost(model, prompt_tokens, completion_tokens)
                    self.total_cost += cost
                
                self._log(f"Model usage: {prompt_tokens} in, {completion_tokens} out")
    
    # ===== Legacy pyautogen API =====
    
    def register_agents(self, agents: List[Any]) -> None:
        """
        Register agents to track their conversations (legacy pyautogen API).
        
        Args:
            agents: List of AutoGen agents (AssistantAgent, UserProxyAgent, etc.)
        """
        if AUTOGEN_VERSION != "legacy":
            self._log("register_agents() is for legacy pyautogen. Using new API.")
            return
        
        self.start_time = time.time()
        
        for agent in agents:
            self._register_single_agent(agent)
            self._registered_agents.append(agent)
        
        self._log(f"Registered {len(agents)} agents")
    
    def _register_single_agent(self, agent: Any) -> None:
        """Register hooks on a single agent."""
        agent_name = getattr(agent, 'name', str(agent))

        # Idempotency guard: skip if already registered to prevent callback stacking
        if agent_name in self._original_callbacks:
            return

        # Hook into receive method
        original_receive = agent.receive
        
        async def tracked_receive(message, sender, request_reply=None, silent=False):
            self._on_message_received(agent_name, message, sender)
            return await original_receive(message, sender, request_reply, silent)
        
        # Check if receive is async
        import asyncio
        if asyncio.iscoroutinefunction(original_receive):
            agent.receive = tracked_receive
        else:
            def sync_tracked_receive(message, sender, request_reply=None, silent=False):
                self._on_message_received(agent_name, message, sender)
                return original_receive(message, sender, request_reply, silent)
            agent.receive = sync_tracked_receive
        
        self._original_callbacks[agent_name] = original_receive
        
        # Hook into function calls if the agent has them
        if hasattr(agent, '_function_map') and agent._function_map:
            self._wrap_functions(agent)
    
    def _wrap_functions(self, agent: Any) -> None:
        """Wrap agent's registered functions for tracking."""
        agent_name = getattr(agent, 'name', str(agent))
        
        for func_name, func in list(agent._function_map.items()):
            original_func = func
            
            def make_wrapper(fn, name):
                def wrapper(*args, **kwargs):
                    self.tool_calls += 1
                    start = time.time()
                    
                    try:
                        result = fn(*args, **kwargs)
                        duration = int((time.time() - start) * 1000)
                        
                        self.client.track_tool_call(
                            session_id=self.session_id,
                            tool_name=name,
                            tool_input={"args": str(args), "kwargs": str(kwargs)},
                            tool_output={"result": str(result)[:500]},
                        )
                        
                        self._log(f"Function {name} completed in {duration}ms")
                        return result
                    except Exception as e:
                        self.client.track_tool_call(
                            session_id=self.session_id,
                            tool_name=name,
                            tool_input={"args": str(args), "kwargs": str(kwargs)},
                            tool_output={"error": str(e)},
                        )
                        raise
                
                return wrapper
            
            agent._function_map[func_name] = make_wrapper(original_func, func_name)
    
    def _on_message_received(self, agent_name: str, message: Any, sender: Any) -> None:
        """Handle a received message (legacy API)."""
        self.message_count += 1
        sender_name = getattr(sender, 'name', str(sender))
        
        # Extract message content
        if isinstance(message, dict):
            content = message.get('content', str(message))
            # Check for function calls
            if 'function_call' in message:
                func_call = message['function_call']
                self._log(f"Function call detected: {func_call.get('name', 'unknown')}")
        else:
            content = str(message)
        
        # Track as user input if first message
        if self.user_input is None and self.message_count == 1:
            self.user_input = content[:500]
        
        # Update assistant output to latest
        self.assistant_output = content[:500]
        
        self._log(f"{sender_name} -> {agent_name}: {content[:50]}...")
        
        self.client.track_decision(
            session_id=self.session_id,
            reasoning=f"[{sender_name} -> {agent_name}] {content[:200]}",
            confidence=1.0
        )
    
    # ===== Session Completion =====
    
    def finish(
        self,
        success: bool = True,
        failure_reason: Optional[str] = None,
        custom_metrics: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Finish the session and record final metrics.
        
        Args:
            success: Whether the agent run succeeded
            failure_reason: Reason for failure if success=False
            custom_metrics: Additional custom metrics
            
        Returns:
            Updated session data
        """
        if self.end_time is None:
            self.end_time = time.time()
        
        # Add AutoGen-specific metrics
        metrics = custom_metrics or {}
        metrics.update({
            "autogen_messages": self.message_count,
            "autogen_tool_calls": self.tool_calls,
            "autogen_version": AUTOGEN_VERSION,
        })
        
        return self.client.complete_session(
            session_id=self.session_id,
            success=success,
            failure_reason=failure_reason,
            estimated_cost=self.total_cost,
            custom_metrics=metrics,
            duration_ms=self.duration_ms,
            prompt_tokens=self.total_prompt_tokens,
            completion_tokens=self.total_completion_tokens,
            total_tokens=self.total_tokens,
            user_input=self.user_input,
            assistant_output=self.assistant_output,
        )
    
    def get_usage_summary(self) -> Dict[str, Any]:
        """Get summary of usage for this session."""
        return {
            "messages": self.message_count,
            "tool_calls": self.tool_calls,
            "total_prompt_tokens": self.total_prompt_tokens,
            "total_completion_tokens": self.total_completion_tokens,
            "total_tokens": self.total_tokens,
            "total_cost": self.total_cost,
            "model": self.model_name,
            "duration_ms": self.duration_ms,
            "autogen_version": AUTOGEN_VERSION,
        }
