"""
Agent Status integration for Microsoft AutoGen.

Provides automatic monitoring for AutoGen agents and conversations.

Usage:

    # Option 1: Wrapper for tracking
    from autogen import AssistantAgent, UserProxyAgent
    from agent_status.integrations.autogen import AgentStatusAgentWrapper, AgentStatusConversationTracker

    assistant = AssistantAgent(name="assistant", llm_config={...})
    wrapped = AgentStatusAgentWrapper(
        agent=assistant,
        api_key="fb_live_xxx",
    )

    user_proxy.initiate_chat(wrapped.agent, message="Hello")

    print(wrapped.metrics)

    # Option 2: Conversation tracker
    tracker = AgentStatusConversationTracker(api_key="fb_live_xxx")
    tracker.start()

    # ... run your AutoGen conversations ...

    tracker.stop()
    print(tracker.metrics)

    # Option 3: Expose for outside-in monitoring
    from agent_status.integrations.autogen import expose

    expose(assistant, api_key="fb_live_xxx", agent_name="My AutoGen Assistant")
"""

from __future__ import annotations

import time
import logging
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

AUTOGEN_AVAILABLE = False

try:
    from autogen_agentchat.agents import AssistantAgent
    ConversableAgent = AssistantAgent
    AUTOGEN_AVAILABLE = True
except ImportError:
    try:
        from autogen import ConversableAgent
        AUTOGEN_AVAILABLE = True
    except ImportError:
        ConversableAgent = object
        AssistantAgent = object

from agent_status.client import AgentStatusClient

logger = logging.getLogger(__name__)


class AgentStatusAgentWrapper:
    """
    Wrapper that adds Agent Status monitoring to an AutoGen agent.

    Tracks:
    - Message counts (sent/received)
    - Response latencies
    - Conversation turns
    - Tool usage

    Args:
        agent: The AutoGen agent to wrap
        api_key: Your API key
        agent_name: Name for this agent in the dashboard
        base_url: API base URL

    Example:
        from autogen import AssistantAgent
        from agent_status.integrations.autogen import AgentStatusAgentWrapper

        assistant = AssistantAgent(name="assistant", llm_config={...})
        wrapped = AgentStatusAgentWrapper(agent=assistant, api_key="fb_live_xxx")

        user_proxy.initiate_chat(wrapped.agent, message="Hello")
        print(wrapped.metrics)
    """

    def __init__(
        self,
        agent: Any,
        api_key: str,
        agent_name: Optional[str] = None,
        base_url: str = "https://api.fabric.carmel.so",
    ):
        if not AUTOGEN_AVAILABLE:
            raise ImportError(
                "AutoGen is not installed. Install with: pip install pyautogen"
            )

        self._original_agent = agent
        self._agent_name = agent_name or getattr(agent, "name", "AutoGen Agent")
        self.api_key = api_key
        self.base_url = base_url

        self._client = AgentStatusClient(api_key=api_key, base_url=base_url)

        self._metrics: Dict[str, Any] = {
            "messages_sent": 0,
            "messages_received": 0,
            "total_latency_ms": 0,
            "conversations": 0,
            "tool_calls": 0,
            "errors": 0,
        }

        self._last_message_time: Optional[float] = None
        self._wrap_agent()

    @property
    def agent(self) -> Any:
        """Get the wrapped agent."""
        return self._original_agent

    @property
    def metrics(self) -> Dict[str, Any]:
        """Get current metrics summary."""
        total_msgs = self._metrics["messages_sent"] + self._metrics["messages_received"]
        return {
            **self._metrics,
            "avg_latency_ms": (
                self._metrics["total_latency_ms"] / self._metrics["messages_received"]
                if self._metrics["messages_received"] > 0 else 0
            ),
            "total_messages": total_msgs,
        }

    def _wrap_agent(self) -> None:
        """Wrap the agent's send/receive methods to track metrics."""
        if hasattr(self._original_agent, "send"):
            original_send = self._original_agent.send

            def wrapped_send(*args, **kwargs):
                self._on_message_sent()
                return original_send(*args, **kwargs)

            self._original_agent.send = wrapped_send

        if hasattr(self._original_agent, "receive"):
            original_receive = self._original_agent.receive

            def wrapped_receive(*args, **kwargs):
                self._on_message_received()
                return original_receive(*args, **kwargs)

            self._original_agent.receive = wrapped_receive

        if hasattr(self._original_agent, "generate_reply"):
            original_generate = self._original_agent.generate_reply

            def wrapped_generate(*args, **kwargs):
                start = time.time()
                try:
                    result = original_generate(*args, **kwargs)
                    latency_ms = (time.time() - start) * 1000
                    self._metrics["total_latency_ms"] += latency_ms
                    return result
                except Exception as e:
                    self._metrics["errors"] += 1
                    raise

            self._original_agent.generate_reply = wrapped_generate

    def _on_message_sent(self) -> None:
        """Called when the agent sends a message."""
        self._metrics["messages_sent"] += 1
        self._last_message_time = time.time()
        logger.debug(f"Agent {self._agent_name} sent message")

    def _on_message_received(self) -> None:
        """Called when the agent receives a message."""
        self._metrics["messages_received"] += 1

        if self._last_message_time:
            latency_ms = (time.time() - self._last_message_time) * 1000
            self._metrics["total_latency_ms"] += latency_ms

        logger.debug(f"Agent {self._agent_name} received message")

    def reset_metrics(self) -> None:
        """Reset all tracked metrics."""
        self._metrics = {
            "messages_sent": 0,
            "messages_received": 0,
            "total_latency_ms": 0,
            "conversations": 0,
            "tool_calls": 0,
            "errors": 0,
        }

    def report(self) -> Dict[str, Any]:
        """Generate a summary report of metrics."""
        return {
            "agent_name": self._agent_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "metrics": self.metrics,
        }


# Backward compatibility alias
RoraAgentWrapper = AgentStatusAgentWrapper


class AgentStatusConversationTracker:
    """
    Tracks metrics across multiple AutoGen conversations.

    This is useful when you have multiple agents and want to track
    overall conversation metrics rather than per-agent.

    Args:
        api_key: Your API key
        session_name: Name for this tracking session
        base_url: API base URL

    Example:
        from agent_status.integrations.autogen import AgentStatusConversationTracker

        tracker = AgentStatusConversationTracker(api_key="fb_live_xxx")
        tracker.start()

        user_proxy.initiate_chat(assistant, message="Hello")

        tracker.stop()
        print(tracker.metrics)
    """

    def __init__(
        self,
        api_key: str,
        session_name: str = "AutoGen Session",
        base_url: str = "https://api.fabric.carmel.so",
    ):
        self.api_key = api_key
        self.session_name = session_name
        self.base_url = base_url

        self._client = AgentStatusClient(api_key=api_key, base_url=base_url)
        self._start_time: Optional[float] = None
        self._end_time: Optional[float] = None
        self._is_tracking = False

        self._metrics: Dict[str, Any] = {
            "total_duration_ms": 0,
            "conversations_tracked": 0,
        }

    @property
    def metrics(self) -> Dict[str, Any]:
        """Get current metrics summary."""
        return {
            **self._metrics,
            "session_name": self.session_name,
            "is_tracking": self._is_tracking,
        }

    def start(self) -> None:
        """Start tracking conversations."""
        self._start_time = time.time()
        self._is_tracking = True
        self._metrics["conversations_tracked"] += 1
        logger.info(f"Started tracking session: {self.session_name}")

    def stop(self) -> None:
        """Stop tracking conversations."""
        if self._start_time:
            self._end_time = time.time()
            duration_ms = (self._end_time - self._start_time) * 1000
            self._metrics["total_duration_ms"] += duration_ms

        self._is_tracking = False
        logger.info(f"Stopped tracking session: {self.session_name}")

    def reset(self) -> None:
        """Reset the tracker."""
        self._start_time = None
        self._end_time = None
        self._is_tracking = False
        self._metrics = {
            "total_duration_ms": 0,
            "conversations_tracked": 0,
        }

    def report(self) -> Dict[str, Any]:
        """Generate a summary report."""
        return {
            "session_name": self.session_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "metrics": self.metrics,
        }


# Backward compatibility alias
RoraConversationTracker = AgentStatusConversationTracker


def expose(
    agent: Any,
    api_key: str,
    agent_name: Optional[str] = None,
    *,
    input_key: str = "input",
    tunnel_url: str = "wss://rora-tunnel.carmel.so/ws/connect",
    blocking: bool = True,
) -> Optional[str]:
    """
    Expose an AutoGen agent for outside-in monitoring.

    This connects your local AutoGen agent to the distributed monitoring network,
    allowing nodes around the world to probe your agent.

    Args:
        agent: An AutoGen ConversableAgent (AssistantAgent, UserProxyAgent, etc.)
        api_key: Your API key
        agent_name: Name for this agent in the dashboard
        input_key: The input key to use when invoking (default: "input")
        tunnel_url: WebSocket URL for the tunnel relay
        blocking: If True, blocks until interrupted. If False, runs in background.

    Returns:
        The public URL where the agent can be reached (if non-blocking)

    Example:
        from autogen import AssistantAgent
        from agent_status.integrations.autogen import expose

        assistant = AssistantAgent(name="assistant", llm_config={...})
        expose(assistant, api_key="fb_live_xxx", agent_name="My Assistant")
    """
    if not AUTOGEN_AVAILABLE:
        raise ImportError(
            "AutoGen is not installed. Install with: pip install pyautogen"
        )

    from agent_status.integrations.langchain import _TunnelHandler

    name = agent_name or getattr(agent, "name", "AutoGen Agent")

    class AutoGenWrapper:
        def __init__(self, agent):
            self.agent = agent

        def invoke(self, inputs):
            input_value = inputs.get(input_key, "")
            if hasattr(self.agent, "generate_reply"):
                result = self.agent.generate_reply(
                    messages=[{"role": "user", "content": input_value}]
                )
                return {"output": str(result) if result else ""}
            elif hasattr(self.agent, "run"):
                result = self.agent.run(input_value)
                return {"output": str(result) if result else ""}
            else:
                return {"output": "Agent does not support direct invocation"}

    handler = _TunnelHandler(
        chain=AutoGenWrapper(agent),
        api_key=api_key,
        agent_name=name,
        input_key=input_key,
        tunnel_url=tunnel_url,
        auto_register=True,
    )

    if blocking:
        handler.run()
        return None
    else:
        import threading
        thread = threading.Thread(target=handler.run, daemon=True)
        thread.start()
        return handler.public_url
