"""AutoGen adapter — instrument your agents in 2 lines.

Usage:
    from velenai_sdk import VelenAI
    from velenai_sdk.adapters.autogen import VelenAIAutoGenHandler

    vai = VelenAI()

    # Option 1: Register as a hook on AssistantAgent
    handler = VelenAIAutoGenHandler(vai)
    assistant = AssistantAgent("analyst", ...)
    assistant.register_hook("process_message_before_send", handler.before_send)

    # Option 2: Wrap an agent
    from velenai_sdk.adapters.autogen import instrument_agent
    assistant = instrument_agent(vai, AssistantAgent("analyst", ...))
"""

from __future__ import annotations

import time
from typing import Any, Callable

from ..client import VelenAI


class VelenAIAutoGenHandler:
    """AutoGen event handler that emits telemetry to VelenAI.

    Can be used as:
    - A hook registered via `register_hook()`
    - A wrapper around `initiate_chat()` / `generate_reply()`
    """

    def __init__(self, client: VelenAI, agent_prefix: str = "autogen"):
        self._client = client
        self._prefix = agent_prefix
        self._chat_starts: dict[str, float] = {}

    def on_chat_start(self, sender: Any, recipient: Any, **kwargs: Any) -> None:
        """Call at the start of a chat/conversation."""
        key = f"{id(sender)}-{id(recipient)}"
        self._chat_starts[key] = time.monotonic()

    def on_chat_end(
        self, sender: Any, recipient: Any, success: bool = True,
        tokens: int | None = None, **kwargs: Any,
    ) -> None:
        """Call at the end of a chat/conversation."""
        key = f"{id(sender)}-{id(recipient)}"
        start = self._chat_starts.pop(key, time.monotonic())
        latency_ms = int((time.monotonic() - start) * 1000)

        agent_name = getattr(sender, "name", "unknown")
        agent_id = f"{self._prefix}.{agent_name}"

        self._client.emit(
            agent_id=agent_id,
            success=success,
            latency_ms=latency_ms,
            tokens=tokens,
            metadata={
                "framework": "autogen",
                "recipient": getattr(recipient, "name", "unknown"),
            },
        )

    def before_send(self, sender: Any, message: Any, recipient: Any, **kwargs: Any) -> Any:
        """Hook for register_hook('process_message_before_send', ...)."""
        key = f"{id(sender)}-{id(recipient)}"
        if key not in self._chat_starts:
            self._chat_starts[key] = time.monotonic()
        return message  # Pass through unmodified

    def after_receive(self, sender: Any, message: Any, recipient: Any, **kwargs: Any) -> Any:
        """Hook for tracking message receipt."""
        key = f"{id(recipient)}-{id(sender)}"
        start = self._chat_starts.pop(key, None)
        if start is not None:
            latency_ms = int((time.monotonic() - start) * 1000)
            agent_name = getattr(recipient, "name", "unknown")
            self._client.emit(
                agent_id=f"{self._prefix}.{agent_name}",
                success=True,
                latency_ms=latency_ms,
                metadata={"framework": "autogen", "type": "message_exchange"},
            )
        return message


def instrument_agent(client: VelenAI, agent: Any, agent_prefix: str = "autogen") -> Any:
    """Wrap an AutoGen agent so all generate_reply calls are tracked.

    Usage:
        assistant = AssistantAgent("analyst", llm_config=...)
        assistant = instrument_agent(vai, assistant)
    """
    original_generate = getattr(agent, "generate_reply", None)
    if original_generate is None:
        return agent

    agent_name = getattr(agent, "name", "unknown")
    agent_id = f"{agent_prefix}.{agent_name}"

    def tracked_generate_reply(*args: Any, **kwargs: Any) -> Any:
        start = time.monotonic()
        success = True
        try:
            result = original_generate(*args, **kwargs)
            return result
        except Exception:
            success = False
            raise
        finally:
            client.emit(
                agent_id=agent_id,
                success=success,
                latency_ms=int((time.monotonic() - start) * 1000),
                metadata={"framework": "autogen", "type": "generate_reply"},
            )

    agent.generate_reply = tracked_generate_reply
    return agent
