"""
AutoGen patches for agent conversation tracing.

Supports two API versions:
- v0.4+ (autogen_agentchat): ChatAgent.on_messages, GroupChatManager
- v0.2.x (autogen): ConversableAgent.generate_reply, ConversableAgent.send

Does NOT suppress provider instrumentation — AutoGen doesn't create its
own LLM spans, so provider patches handle LLM calls naturally.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Dict, TypeVar

import wrapt

from risicare.integrations._base import (
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
    create_framework_attributes,
    get_framework_version,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


def _get_agent_name(instance: Any) -> str:
    """Extract agent name from an AutoGen agent instance."""
    return (
        getattr(instance, "name", None)
        or getattr(instance, "_name", None)
        or type(instance).__name__
    )


# =============================================================================
# v0.2.x Patches (ConversableAgent)
# =============================================================================


def _wrap_generate_reply(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ConversableAgent.generate_reply (v0.2.x)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind, SemanticPhase

    agent_name = _get_agent_name(instance)
    version = get_framework_version("pyautogen") or get_framework_version("autogen")

    attrs = create_framework_attributes(
        "autogen", version, agent_name=agent_name
    )
    attrs["agent.role"] = "worker"
    attrs["agent.type"] = "autogen"
    attrs["agent.name"] = agent_name

    # Capture message context
    messages = args[0] if args else kwargs.get("messages")
    if messages and isinstance(messages, list):
        attrs["framework.autogen.message_count"] = len(messages)
        if should_trace_content() and messages:
            last_msg = messages[-1]
            if isinstance(last_msg, dict):
                content = last_msg.get("content", "")
                if isinstance(content, str):
                    attrs["framework.autogen.last_message"] = truncate_content(
                        content, 2000
                    )
                attrs["framework.autogen.last_sender"] = last_msg.get("name", "")

    sender = args[1] if len(args) > 1 else kwargs.get("sender")
    if sender:
        attrs["framework.autogen.sender"] = _get_agent_name(sender)

    with tracer.start_span(
        name=f"autogen.agent/{agent_name}/reply",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        span.semantic_phase = SemanticPhase.THINK
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and result is not None:
                if isinstance(result, str):
                    safe_set_attribute(
                        span,
                        "framework.autogen.reply",
                        truncate_content(result, 5000),
                    )
                elif isinstance(result, dict):
                    content = result.get("content", "")
                    if isinstance(content, str):
                        safe_set_attribute(
                            span,
                            "framework.autogen.reply",
                            truncate_content(content, 5000),
                        )

            return result
        except Exception as e:
            record_error(span, e)
            raise


def _wrap_send(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ConversableAgent.send (v0.2.x)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind, SemanticPhase

    sender_name = _get_agent_name(instance)
    recipient = args[1] if len(args) > 1 else kwargs.get("recipient")
    recipient_name = _get_agent_name(recipient) if recipient else "unknown"

    attrs = create_framework_attributes(
        "autogen",
        sender=sender_name,
        receiver=recipient_name,
    )

    message = args[0] if args else kwargs.get("message")
    if message and should_trace_content():
        if isinstance(message, str):
            attrs["framework.autogen.message_content"] = truncate_content(
                message, 2000
            )
        elif isinstance(message, dict):
            content = message.get("content", "")
            if isinstance(content, str):
                attrs["framework.autogen.message_content"] = truncate_content(
                    content, 2000
                )

    with tracer.start_span(
        name=f"autogen.message/{sender_name}→{recipient_name}",
        kind=SpanKind.INTERNAL,
        attributes=attrs,
    ) as span:
        span.semantic_phase = SemanticPhase.COMMUNICATE
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)
            return result
        except Exception as e:
            record_error(span, e)
            raise


def _wrap_initiate_chat(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ConversableAgent.initiate_chat (v0.2.x root span)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    initiator_name = _get_agent_name(instance)
    recipient = args[0] if args else kwargs.get("recipient")
    recipient_name = _get_agent_name(recipient) if recipient else "unknown"
    version = get_framework_version("pyautogen") or get_framework_version("autogen")

    attrs = create_framework_attributes(
        "autogen", version,
        initiator=initiator_name,
        recipient=recipient_name,
    )
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "autogen"

    max_turns = kwargs.get("max_turns")
    if max_turns:
        attrs["framework.autogen.max_turns"] = max_turns

    with tracer.start_span(
        name=f"autogen.chat/{initiator_name}→{recipient_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Capture chat summary
            chat_history = getattr(result, "chat_history", None)
            if chat_history:
                span.set_attribute(
                    "framework.autogen.total_messages", len(chat_history)
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


# =============================================================================
# v0.4+ Patches (autogen_agentchat)
# =============================================================================


async def _wrap_on_messages(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for ChatAgent.on_messages (v0.4+)."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind, SemanticPhase

    agent_name = _get_agent_name(instance)
    version = get_framework_version("autogen-agentchat")

    attrs = create_framework_attributes(
        "autogen", version, agent_name=agent_name
    )
    attrs["agent.role"] = "worker"
    attrs["agent.type"] = "autogen"
    attrs["agent.name"] = agent_name

    messages = args[0] if args else kwargs.get("messages", [])
    if messages:
        attrs["framework.autogen.input_message_count"] = len(messages)

    with tracer.start_span(
        name=f"autogen.agent/{agent_name}/turn",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        span.semantic_phase = SemanticPhase.THINK
        try:
            start_time = time.perf_counter()
            result = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Capture response
            response = getattr(result, "chat_message", None)
            if response and should_trace_content():
                content = getattr(response, "content", "")
                if isinstance(content, str):
                    safe_set_attribute(
                        span,
                        "framework.autogen.response",
                        truncate_content(content, 5000),
                    )

            return result
        except Exception as e:
            record_error(span, e)
            raise


# =============================================================================
# Patch Entry Point
# =============================================================================


def patch_autogen(module: Any) -> None:
    """Apply wrapt patches to AutoGen. Handles both v0.2.x and v0.4+."""
    patched = False

    # v0.4+ (autogen_agentchat)
    try:
        wrapt.wrap_function_wrapper(
            "autogen_agentchat.base",
            "ChatAgent.on_messages",
            _wrap_on_messages,
        )
        patched = True
        logger.debug("Patched AutoGen v0.4+ ChatAgent.on_messages")
    except Exception:
        pass

    # v0.2.x (ConversableAgent)
    try:
        wrapt.wrap_function_wrapper(
            module,
            "ConversableAgent.generate_reply",
            _wrap_generate_reply,
        )
        patched = True
        logger.debug("Patched AutoGen v0.2.x ConversableAgent.generate_reply")
    except Exception:
        pass

    try:
        wrapt.wrap_function_wrapper(
            module,
            "ConversableAgent.send",
            _wrap_send,
        )
    except Exception:
        pass

    try:
        wrapt.wrap_function_wrapper(
            module,
            "ConversableAgent.initiate_chat",
            _wrap_initiate_chat,
        )
    except Exception:
        pass

    if not patched:
        logger.debug("No AutoGen classes found to patch")
