"""Auto-instrumentation for popular LLM frameworks using OpenInference."""

import logging
from typing import Optional, List, Dict
from opentelemetry.sdk.trace import TracerProvider

logger = logging.getLogger(__name__)

# Configuration for all supported OpenInference instrumentors
# Easy to extend - just add new entries when OpenInference adds support for new frameworks
INSTRUMENTORS = [
    # LLM Frameworks (priority - most commonly used)
    {
        "name": "openai",
        "package": "openinference.instrumentation.openai",
        "class": "OpenAIInstrumentor",
        "description": "OpenAI SDK"
    },
    {
        "name": "llamaindex",
        "package": "openinference.instrumentation.llamaindex",
        "class": "LlamaIndexInstrumentor",
        "description": "LlamaIndex"
    },
    {
        "name": "langchain",
        "package": "openinference.instrumentation.langchain",
        "class": "LangChainInstrumentor",
        "description": "LangChain"
    },

    # LLM Providers
    {
        "name": "vertexai",
        "package": "openinference.instrumentation.vertexai",
        "class": "VertexAIInstrumentor",
        "description": "Google VertexAI"
    },
    {
        "name": "bedrock",
        "package": "openinference.instrumentation.bedrock",
        "class": "BedrockInstrumentor",
        "description": "AWS Bedrock"
    },
    {
        "name": "anthropic",
        "package": "openinference.instrumentation.anthropic",
        "class": "AnthropicInstrumentor",
        "description": "Anthropic Claude"
    },
    {
        "name": "mistralai",
        "package": "openinference.instrumentation.mistralai",
        "class": "MistralAIInstrumentor",
        "description": "MistralAI"
    },
    {
        "name": "groq",
        "package": "openinference.instrumentation.groq",
        "class": "GroqInstrumentor",
        "description": "Groq"
    },
    {
        "name": "litellm",
        "package": "openinference.instrumentation.litellm",
        "class": "LiteLLMInstrumentor",
        "description": "liteLLM"
    },

    # Other Tools
    {
        "name": "dspy",
        "package": "openinference.instrumentation.dspy",
        "class": "DSPyInstrumentor",
        "description": "DSPy"
    },
    {
        "name": "crewai",
        "package": "openinference.instrumentation.crewai",
        "class": "CrewAIInstrumentor",
        "description": "CrewAI"
    },
    {
        "name": "haystack",
        "package": "openinference.instrumentation.haystack",
        "class": "HaystackInstrumentor",
        "description": "Haystack"
    },
    {
        "name": "instructor",
        "package": "openinference.instrumentation.instructor",
        "class": "InstructorInstrumentor",
        "description": "Instructor"
    },
    {
        "name": "guardrails",
        "package": "openinference.instrumentation.guardrails",
        "class": "GuardrailsInstrumentor",
        "description": "Guardrails"
    },
]


def auto_instrument(
    tracer_provider: Optional[TracerProvider] = None,
    suppress_errors: bool = True,
    include_libraries: Optional[List[str]] = None,
    exclude_libraries: Optional[List[str]] = None,
) -> Dict[str, str]:
    """Automatically instrument all available LLM libraries using OpenInference.

    This function attempts to instrument all OpenInference-supported libraries that are installed.
    Only libraries actually installed are instrumented - missing libraries are silently skipped.

    This is the easiest way to get started with Coalex observability. Just call this function
    after register() and all your LLM library calls will automatically be traced with token counts.

    Args:
        tracer_provider: TracerProvider to use (defaults to global provider)
        suppress_errors: If True, log errors instead of raising (default: True)
        include_libraries: If provided, only instrument these libraries
        exclude_libraries: If provided, skip these libraries

    Returns:
        Dictionary mapping library name to instrumentation status:
        {
            "openai": "success" | "not_installed" | "already_instrumented" | "error" | "excluded",
            "langchain": ...,
            "vertexai": ...,
        }

    Example:
        # Auto-instrument everything after registration
        tracer_provider = coalex.register(agent_id="...")
        results = coalex.auto_instrument(tracer_provider=tracer_provider)
        print(f"Instrumented: {[k for k,v in results.items() if v == 'success']}")

        # Only instrument specific libraries
        coalex.auto_instrument(
            tracer_provider=tracer_provider,
            include_libraries=["openai", "langchain"]
        )
    """
    results = {}

    # Filter based on include/exclude
    for config in INSTRUMENTORS:
        lib_name = config["name"]

        # Check include/exclude filters
        if include_libraries and lib_name not in include_libraries:
            continue
        if exclude_libraries and lib_name in exclude_libraries:
            results[lib_name] = "excluded"
            continue

        results[lib_name] = _instrument_library(
            config=config,
            tracer_provider=tracer_provider,
            suppress_errors=suppress_errors
        )

    # Log summary
    success_count = sum(1 for v in results.values() if v == "success")
    total_attempted = sum(1 for v in results.values() if v != "excluded")

    if success_count > 0:
        successful_libs = [k for k, v in results.items() if v == "success"]
        logger.info(
            f"Auto-instrumentation complete: {success_count}/{total_attempted} libraries instrumented "
            f"({', '.join(successful_libs)})"
        )
    else:
        logger.info(
            "Auto-instrumentation complete: no libraries were instrumented. "
            "Install OpenInference packages to enable auto-instrumentation: "
            "pip install 'coalex[auto-instrument]'"
        )

    logger.debug(f"Full instrumentation results: {results}")

    return results


def _instrument_library(
    config: Dict[str, str],
    tracer_provider: Optional[TracerProvider],
    suppress_errors: bool
) -> str:
    """Instrument a single library.

    Args:
        config: Instrumentor configuration dict with name, package, class, description
        tracer_provider: Optional TracerProvider to use
        suppress_errors: Whether to suppress and log errors

    Returns:
        Status string: "success", "not_installed", "already_instrumented", or "error"
    """
    lib_name = config["name"]
    package_name = config["package"]
    class_name = config["class"]
    description = config["description"]

    try:
        # Try to import the instrumentor
        try:
            module = __import__(package_name, fromlist=[class_name])
            instrumentor_class = getattr(module, class_name)
        except ImportError:
            logger.debug(f"Skipping {lib_name}: {package_name} not installed")
            return "not_installed"
        except AttributeError:
            logger.warning(f"Found {package_name} but missing {class_name} class")
            return "error"

        # Create instrumentor instance
        instrumentor = instrumentor_class()

        # Check if already instrumented
        if hasattr(instrumentor, '_is_instrumented_by_opentelemetry'):
            if instrumentor._is_instrumented_by_opentelemetry:
                logger.debug(f"{lib_name} already instrumented")
                return "already_instrumented"

        # Instrument with the provided tracer provider
        if tracer_provider:
            instrumentor.instrument(tracer_provider=tracer_provider)
        else:
            instrumentor.instrument()

        logger.info(f"✓ Successfully instrumented {description} ({lib_name})")
        return "success"

    except Exception as e:
        error_msg = f"Failed to instrument {lib_name}: {type(e).__name__}: {str(e)}"

        if suppress_errors:
            logger.warning(error_msg)
            return "error"
        else:
            logger.error(error_msg)
            raise


# Convenience functions for commonly used libraries

def instrument_openai(tracer_provider: Optional[TracerProvider] = None) -> bool:
    """Instrument OpenAI SDK specifically.

    Args:
        tracer_provider: Optional TracerProvider to use

    Returns:
        True if instrumentation succeeded, False otherwise
    """
    result = _instrument_library(
        config={
            "name": "openai",
            "package": "openinference.instrumentation.openai",
            "class": "OpenAIInstrumentor",
            "description": "OpenAI SDK"
        },
        tracer_provider=tracer_provider,
        suppress_errors=True
    )
    return result == "success"


def instrument_llamaindex(tracer_provider: Optional[TracerProvider] = None) -> bool:
    """Instrument LlamaIndex specifically.

    Args:
        tracer_provider: Optional TracerProvider to use

    Returns:
        True if instrumentation succeeded, False otherwise
    """
    result = _instrument_library(
        config={
            "name": "llamaindex",
            "package": "openinference.instrumentation.llamaindex",
            "class": "LlamaIndexInstrumentor",
            "description": "LlamaIndex"
        },
        tracer_provider=tracer_provider,
        suppress_errors=True
    )
    return result == "success"


def instrument_langchain(tracer_provider: Optional[TracerProvider] = None) -> bool:
    """Instrument LangChain specifically.

    Args:
        tracer_provider: Optional TracerProvider to use

    Returns:
        True if instrumentation succeeded, False otherwise
    """
    result = _instrument_library(
        config={
            "name": "langchain",
            "package": "openinference.instrumentation.langchain",
            "class": "LangChainInstrumentor",
            "description": "LangChain Framework"
        },
        tracer_provider=tracer_provider,
        suppress_errors=True
    )
    return result == "success"


def instrument_vertexai(tracer_provider: Optional[TracerProvider] = None) -> bool:
    """Instrument Google VertexAI specifically.

    Args:
        tracer_provider: Optional TracerProvider to use

    Returns:
        True if instrumentation succeeded, False otherwise
    """
    result = _instrument_library(
        config={
            "name": "vertexai",
            "package": "openinference.instrumentation.vertexai",
            "class": "VertexAIInstrumentor",
            "description": "Google VertexAI"
        },
        tracer_provider=tracer_provider,
        suppress_errors=True
    )
    return result == "success"


def instrument_bedrock(tracer_provider: Optional[TracerProvider] = None) -> bool:
    """Instrument AWS Bedrock specifically.

    Args:
        tracer_provider: Optional TracerProvider to use

    Returns:
        True if instrumentation succeeded, False otherwise
    """
    result = _instrument_library(
        config={
            "name": "bedrock",
            "package": "openinference.instrumentation.bedrock",
            "class": "BedrockInstrumentor",
            "description": "AWS Bedrock"
        },
        tracer_provider=tracer_provider,
        suppress_errors=True
    )
    return result == "success"


def instrument_anthropic(tracer_provider: Optional[TracerProvider] = None) -> bool:
    """Instrument Anthropic Claude specifically.

    Args:
        tracer_provider: Optional TracerProvider to use

    Returns:
        True if instrumentation succeeded, False otherwise
    """
    result = _instrument_library(
        config={
            "name": "anthropic",
            "package": "openinference.instrumentation.anthropic",
            "class": "AnthropicInstrumentor",
            "description": "Anthropic Claude"
        },
        tracer_provider=tracer_provider,
        suppress_errors=True
    )
    return result == "success"
