"""
Agno auto-instrumentation.

Automatically patches Agno Agent and Team initialization to inject Aigie tracing.
"""

import functools
import logging
from typing import Set

logger = logging.getLogger(__name__)

_patched_classes: Set[int] = set()
_original_agent_init = None
_original_team_init = None


def patch_agno() -> bool:
    """
    Patch Agno for auto-instrumentation.

    This patches Agent.__init__ and Team.__init__ to automatically register
    AgnoHandler when agents/teams are created.

    Returns:
        True if patching was successful (or already patched), False if agno is not installed.
    """
    global _original_agent_init, _original_team_init

    try:
        from agno.agent import Agent
    except ImportError:
        logger.debug("[AIGIE] Agno not installed, skipping auto-instrumentation")
        return False

    patched_any = False

    # Patch Agent
    if id(Agent.__init__) not in _patched_classes:
        _original_agent_init = Agent.__init__
        orig_agent_init = _original_agent_init

        @functools.wraps(orig_agent_init)
        def patched_agent_init(self, *args, **kwargs):
            """Patched Agent.__init__ that auto-registers Aigie handler."""
            # Call original __init__ first so the agent is fully constructed
            orig_agent_init(self, *args, **kwargs)

            try:
                from .config import AgnoConfig
                from .handler import AgnoHandler

                config = AgnoConfig.from_env()
                if config.enabled:
                    # Check if a handler is already registered
                    pre_hooks = getattr(self, "pre_hooks", None) or []
                    has_aigie = any(
                        getattr(h, "_aigie_handler", None) is not None for h in pre_hooks
                    )
                    if not has_aigie:
                        handler = AgnoHandler(config=config)
                        handler.register(self)
                        logger.debug("[AIGIE] Auto-registered AgnoHandler on Agent")
            except Exception as e:
                logger.warning(f"[AIGIE] Failed to auto-register AgnoHandler on Agent: {e}")

        Agent.__init__ = patched_agent_init
        _patched_classes.add(id(Agent.__init__))
        patched_any = True
        logger.debug("[AIGIE] Agno Agent patched for auto-instrumentation")

    # Patch Team (optional — may not be available in all agno versions)
    try:
        from agno.team import Team

        if id(Team.__init__) not in _patched_classes:
            _original_team_init = Team.__init__
            orig_team_init = _original_team_init

            @functools.wraps(orig_team_init)
            def patched_team_init(self, *args, **kwargs):
                """Patched Team.__init__ that auto-registers Aigie handler."""
                orig_team_init(self, *args, **kwargs)

                try:
                    from .config import AgnoConfig
                    from .handler import AgnoHandler

                    config = AgnoConfig.from_env()
                    if config.enabled and config.trace_teams:
                        pre_hooks = getattr(self, "pre_hooks", None) or []
                        has_aigie = any(
                            getattr(h, "_aigie_handler", None) is not None for h in pre_hooks
                        )
                        if not has_aigie:
                            handler = AgnoHandler(config=config)
                            handler.register(self)
                            logger.debug("[AIGIE] Auto-registered AgnoHandler on Team")
                except Exception as e:
                    logger.warning(f"[AIGIE] Failed to auto-register AgnoHandler on Team: {e}")

            Team.__init__ = patched_team_init
            _patched_classes.add(id(Team.__init__))
            patched_any = True
            logger.debug("[AIGIE] Agno Team patched for auto-instrumentation")

    except ImportError:
        logger.debug("[AIGIE] Agno Team not available, skipping Team patch")

    if patched_any:
        logger.info("[AIGIE] Agno patched for auto-instrumentation")

    return patched_any


def unpatch_agno() -> None:
    """Remove Agno patches (for testing)."""
    global _original_agent_init, _original_team_init

    try:
        from agno.agent import Agent

        if _original_agent_init is not None:
            Agent.__init__ = _original_agent_init
            _original_agent_init = None
    except ImportError:
        pass

    try:
        from agno.team import Team

        if _original_team_init is not None:
            Team.__init__ = _original_team_init
            _original_team_init = None
    except ImportError:
        pass

    _patched_classes.clear()
    logger.info("[AIGIE] Agno patches removed")


def is_agno_patched() -> bool:
    """
    Check if Agno is currently patched.

    Returns:
        True if patched, False otherwise.
    """
    try:
        from agno.agent import Agent

        return id(Agent.__init__) in _patched_classes
    except ImportError:
        return False
