# Copyright (c) Microsoft. All rights reserved.

"""
Instrumentation module for mantisdk.

Provides automatic metric collection and tracing for various libraries.
Uses a plugin architecture with an InstrumentorRegistry.

Usage:
    # Enable specific instrumentors
    mantisdk.instrument("openai")
    mantisdk.instrument("openai", "anthropic", "langchain")

    # Enable all available instrumentors
    mantisdk.instrument_all()

    # Disable instrumentors
    mantisdk.uninstrument("openai")
    mantisdk.uninstrument_all()
"""

import warnings
from typing import List, Optional

# Export base classes and registry
from .base import Instrumentor, MetricSchema
from .registry import (
    InstrumentorRegistry,
)
from .registry import instrument as registry_instrument
from .registry import uninstrument as registry_uninstrument

__all__ = [
    "Instrumentor",
    "MetricSchema",
    "InstrumentorRegistry",
    "instrument",
    "uninstrument",
    "instrument_all",
    "uninstrument_all",
]

# Legacy instrumentor availability flags
AGENTOPS_INSTALLED: bool = False
AGENTOPS_LANGCHAIN_INSTALLED: bool = False
LITELLM_INSTALLED: bool = False
VLLM_INSTALLED: bool = False
WEAVE_INSTALLED: bool = False
OPENAI_OTEL_INSTALLED: bool = False

try:
    from . import agentops  # type: ignore

    AGENTOPS_INSTALLED = True  # type: ignore
except ImportError:
    pass

try:
    from . import litellm  # type: ignore

    LITELLM_INSTALLED = True  # type: ignore
except ImportError:
    pass

try:
    from opentelemetry.instrumentation.openai import OpenAIInstrumentor  # type: ignore

    OPENAI_OTEL_INSTALLED = True
except ImportError:
    pass


def instrument(*names: str, **kwargs) -> List[Optional[Instrumentor]]:
    """
    Activate one or more instrumentors by name.

    Uses the InstrumentorRegistry to activate instrumentors.
    For legacy instrumentors not yet migrated to the registry,
    falls back to direct activation.

    Args:
        *names: Instrumentor names to activate
        **kwargs: Additional arguments passed to instrumentor constructors

    Returns:
        List of activated instrumentor instances (None for failed activations)

    Examples:
        # Enable OpenAI instrumentation
        mantisdk.instrument("openai")

        # Enable multiple instrumentors
        mantisdk.instrument("openai", "anthropic", "langchain")

        # Enable all available instrumentors
        mantisdk.instrument_all()
    """
    results = []
    for name in names:
        # Try registry first
        try:
            result = InstrumentorRegistry.instrument(name, **kwargs)
            results.append(result)
        except ValueError:
            # Fall back to legacy instrumentors
            result = _instrument_legacy(name)
            results.append(result)
    return results


def uninstrument(*names: str) -> None:
    """
    Deactivate one or more instrumentors by name.

    Args:
        *names: Instrumentor names to deactivate
    """
    for name in names:
        # Try registry first
        if InstrumentorRegistry.is_active(name):
            InstrumentorRegistry.uninstrument(name)
        else:
            # Fall back to legacy uninstrument
            _uninstrument_legacy(name)


def _instrument_legacy(name: str) -> Optional[Instrumentor]:
    """Activate a legacy instrumentor by name."""
    name_lower = name.lower()

    if name_lower == "agentops" and AGENTOPS_INSTALLED:
        from .agentops import instrument_agentops

        instrument_agentops()
        return None  # Legacy doesn't return instance

    if name_lower == "litellm" and LITELLM_INSTALLED:
        from .litellm import instrument_litellm

        instrument_litellm()
        return None

    if name_lower == "vllm" and VLLM_INSTALLED:
        from .vllm import instrument_vllm

        instrument_vllm()
        return None

    if name_lower in ("openai", "openai-otel") and OPENAI_OTEL_INSTALLED:
        try:
            OpenAIInstrumentor().instrument()  # type: ignore
        except Exception as e:
            warnings.warn(f"Failed to instrument OpenAI SDK: {e}")
        return None

    warnings.warn(f"Unknown or unavailable instrumentor: {name}")
    return None


def _uninstrument_legacy(name: str) -> None:
    """Deactivate a legacy instrumentor by name."""
    name_lower = name.lower()

    if name_lower == "agentops" and AGENTOPS_INSTALLED:
        try:
            from .agentops import uninstrument_agentops

            uninstrument_agentops()
        except ImportError:
            pass

    elif name_lower == "litellm" and LITELLM_INSTALLED:
        try:
            from .litellm import uninstrument_litellm

            uninstrument_litellm()
        except ImportError:
            pass

    elif name_lower == "vllm" and VLLM_INSTALLED:
        try:
            from .vllm import uninstrument_vllm

            uninstrument_vllm()
        except ImportError:
            pass

    elif name_lower in ("openai", "openai-otel") and OPENAI_OTEL_INSTALLED:
        try:
            OpenAIInstrumentor().uninstrument()  # type: ignore
        except Exception:
            pass


def instrument_all():
    """
    Instrument all available instrumentation libraries.

    This includes both registry-based instrumentors and legacy instrumentors.
    """
    # First, activate all registry-based instrumentors
    for name in InstrumentorRegistry.list_available():
        try:
            InstrumentorRegistry.instrument(name)
        except Exception as e:
            warnings.warn(f"Failed to instrument {name}: {e}")

    # Then activate legacy instrumentors
    if AGENTOPS_INSTALLED:
        from .agentops import instrument_agentops

        instrument_agentops()
    else:
        warnings.warn("agentops is not installed. It's therefore not instrumented.")

    if LITELLM_INSTALLED:
        from .litellm import instrument_litellm

        instrument_litellm()
    else:
        warnings.warn("litellm is not installed. It's therefore not instrumented.")

    if VLLM_INSTALLED:
        from .vllm import instrument_vllm

        instrument_vllm()
    else:
        warnings.warn("vllm is not installed. It's therefore not instrumented.")

    if AGENTOPS_LANGCHAIN_INSTALLED:
        from .agentops_langchain import instrument_agentops_langchain

        instrument_agentops_langchain()
    else:
        warnings.warn("Agentops-langchain integration is not installed. It's therefore not instrumented.")

    # Enable standard OpenAI instrumentation if available
    # This allows client-side spans to be generated for OpenAI SDK calls
    if OPENAI_OTEL_INSTALLED:
        try:
            OpenAIInstrumentor().instrument()  # type: ignore
        except Exception as e:
            warnings.warn(f"Failed to instrument OpenAI SDK: {e}")


def uninstrument_all():
    """
    Uninstrument all active instrumentation libraries.

    This includes both registry-based instrumentors and legacy instrumentors.
    """
    # First, deactivate all registry-based instrumentors
    InstrumentorRegistry.uninstrument_all()

    # Then deactivate legacy instrumentors
    if AGENTOPS_INSTALLED:
        try:
            from .agentops import uninstrument_agentops

            uninstrument_agentops()
        except ImportError:
            warnings.warn("agentops is installed but uninstrument_agentops could not be imported.")
    else:
        warnings.warn("agentops is not installed. It's therefore not uninstrumented.")

    if LITELLM_INSTALLED:
        try:
            from .litellm import uninstrument_litellm

            uninstrument_litellm()
        except ImportError:
            warnings.warn("litellm is installed but uninstrument_litellm could not be imported.")
    else:
        warnings.warn("litellm is not installed. It's therefore not uninstrumented.")

    if VLLM_INSTALLED:
        try:
            from .vllm import uninstrument_vllm

            uninstrument_vllm()
        except ImportError:
            warnings.warn("vllm is installed but uninstrument_vllm could not be imported.")
    else:
        warnings.warn("vllm is not installed. It's therefore not uninstrumented.")

    if AGENTOPS_LANGCHAIN_INSTALLED:
        try:
            from .agentops_langchain import uninstrument_agentops_langchain

            uninstrument_agentops_langchain()
        except ImportError:
            warnings.warn("agentops_langchain is installed but uninstrument_agentops_langchain could not be imported.")
    else:
        warnings.warn("Agentops-langchain integration is not installed. It's therefore not uninstrumented.")

    if OPENAI_OTEL_INSTALLED:
        try:
            OpenAIInstrumentor().uninstrument()  # type: ignore
        except Exception:
            pass
