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

import logging

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: TracerProvider | None = None,
    suppress_errors: bool = True,
    include_libraries: list[str] | None = None,
    exclude_libraries: list[str] | None = None,
) -> dict[str, str]:
    """Automatically instrument all available LLM libraries using OpenInference.

    Only libraries actually installed are instrumented — missing libraries
    are silently skipped.

    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
        ("success", "not_installed", "already_instrumented", "error", "excluded").
    """
    results: dict[str, str] = {}

    for config in INSTRUMENTORS:
        lib_name = config["name"]

        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,
        )

    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: TracerProvider | None,
    suppress_errors: bool,
) -> str:
    """Instrument a single library.

    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:
            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"

        instrumentor = instrumentor_class()

        is_instrumented = getattr(instrumentor, "_is_instrumented_by_opentelemetry", False)
        if is_instrumented:
            logger.debug(f"{lib_name} already instrumented")
            return "already_instrumented"

        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__}: {e}"

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


# Convenience functions for commonly used libraries


def _convenience_instrument(name: str, tracer_provider: TracerProvider | None) -> bool:
    """Helper for convenience functions — find config by name and instrument."""
    config = next((c for c in INSTRUMENTORS if c["name"] == name), None)
    if not config:
        return False
    return _instrument_library(config=config, tracer_provider=tracer_provider, suppress_errors=True) == "success"


def instrument_openai(tracer_provider: TracerProvider | None = None) -> bool:
    """Instrument OpenAI SDK specifically."""
    return _convenience_instrument("openai", tracer_provider)


def instrument_llamaindex(tracer_provider: TracerProvider | None = None) -> bool:
    """Instrument LlamaIndex specifically."""
    return _convenience_instrument("llamaindex", tracer_provider)


def instrument_langchain(tracer_provider: TracerProvider | None = None) -> bool:
    """Instrument LangChain specifically."""
    return _convenience_instrument("langchain", tracer_provider)


def instrument_vertexai(tracer_provider: TracerProvider | None = None) -> bool:
    """Instrument Google VertexAI specifically."""
    return _convenience_instrument("vertexai", tracer_provider)


def instrument_bedrock(tracer_provider: TracerProvider | None = None) -> bool:
    """Instrument AWS Bedrock specifically."""
    return _convenience_instrument("bedrock", tracer_provider)


def instrument_anthropic(tracer_provider: TracerProvider | None = None) -> bool:
    """Instrument Anthropic Claude specifically."""
    return _convenience_instrument("anthropic", tracer_provider)
