from __future__ import annotations

import importlib
import inspect
import logging
from collections.abc import Callable
from opentelemetry.instrumentation.dependencies import DependencyConflictError

logger = logging.getLogger("impact.sdk.registry")


def _supports_raise_exception_flag(instrument: Callable[..., object]) -> bool:
    """Return True if the instrument callable supports raise_exception_on_conflict."""
    signature = inspect.signature(instrument)
    parameters = signature.parameters
    if "raise_exception_on_conflict" in parameters:
        return True
    return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values())


def _instrument(module: str, class_name: str) -> bool:
    """Try to import `module`, get `class_name`, and call `.instrument()` (best-effort, no hard failures)."""
    try:
        mod = importlib.import_module(module)
        cls = getattr(mod, class_name)
        inst = cls()
        instrument_kwargs: dict[str, object] = {}
        instrument_method = inst.instrument
        if _supports_raise_exception_flag(instrument_method):
            # Never raise on conflicts — don't break customer code.
            instrument_kwargs["raise_exception_on_conflict"] = False
        instrument_method(**instrument_kwargs)
        return True
    except ModuleNotFoundError as exc:
        # Normal case: user hasn't installed that provider/instrumentation.
        logger.debug("Skipping instrumentation for %s: dependency %s is not installed", module, exc.name)
        return False
    except DependencyConflictError as exc:
        # Another vendor may have pinned versions; skip rather than fail hard.
        logger.warning("Skipping instrumentation for %s due to dependency conflict: %s", module, exc)
        return False
    except Exception:
        logger.exception("Instrumentation for %s.%s failed", module, class_name)
        return False


# --- Open Telemetry ---

def instrument_google_genai() -> bool:
    return _instrument("opentelemetry.instrumentation.google_genai", "GoogleGenAiSdkInstrumentor")

def instrument_openai_agents() -> bool:
    return _instrument("opentelemetry.instrumentation.openai_agents", "OpenAIAgentsInstrumentor")

# --- Custom Implementation ---

def instrument_foundry_agents() -> bool:
    """Instrument AIAgentsInstrumentor. Requires azure-core-tracing-opentelemetry to be installed."""
    try:
        from azure.ai.agents.telemetry import AIAgentsInstrumentor  # type: ignore[import]
        instrumentor = AIAgentsInstrumentor()
        if not instrumentor.is_instrumented():
            instrumentor.instrument(enable_content_recording=True)
            logger.debug("Foundry Agents: instrumented AgentsClient class")
        return True
    except (ImportError, ModuleNotFoundError):
        logger.debug("Skipping Foundry Agents: azure-ai-agents or azure-core-tracing-opentelemetry not installed")
        return False
    except Exception:
        logger.exception("Foundry Agents instrumentation failed")
        return False

def instrument_agent_framework() -> bool:
    """Instrument Microsoft Agent Framework.
    Agent Framework self-instruments via class decorators but gates span emission
    behind an enable_instrumentation flag (default False). We call enable_instrumentation()
    directly to ensure it works regardless of import ordering (the env var alone may be
    read too early if the customer imports agent_framework before calling impact.init()).
    """
    try:
        from agent_framework.observability import enable_instrumentation  # type: ignore[import]
        enable_instrumentation()
        logger.debug("Microsoft Agent Framework: enabled instrumentation")
        return True
    except ImportError:
        logger.debug("Skipping Agent Framework: agent-framework is not installed")
        return False
    except Exception:
        logger.exception("Agent Framework instrumentation failed")
        return False

# --- OpenLLMetry ---

def instrument_openai() -> bool:
    return _instrument("opentelemetry.instrumentation.openai", "OpenAIInstrumentor")

def instrument_aws_bedrock() -> bool:
    return _instrument("opentelemetry.instrumentation.bedrock", "BedrockInstrumentor")

def instrument_anthropic() -> bool:
    return _instrument("opentelemetry.instrumentation.anthropic", "AnthropicInstrumentor")

def instrument_cohere() -> bool:
    return _instrument("opentelemetry.instrumentation.cohere", "CohereInstrumentor")

def instrument_mistral() -> bool:
    return _instrument("opentelemetry.instrumentation.mistralai", "MistralAiInstrumentor")

def instrument_watsonx() -> bool:
    return _instrument("opentelemetry.instrumentation.watsonx", "WatsonxInstrumentor")

def instrument_groq() -> bool:
    return _instrument("opentelemetry.instrumentation.groq", "GroqInstrumentor")

def instrument_together() -> bool:
    return _instrument("opentelemetry.instrumentation.together", "TogetherInstrumentor")

def instrument_ollama() -> bool:
    return _instrument("opentelemetry.instrumentation.ollama", "OllamaInstrumentor")

def instrument_langchain() -> bool:
    return _instrument("opentelemetry.instrumentation.langchain", "LangchainInstrumentor")

def instrument_llama_index() -> bool:
    return _instrument("opentelemetry.instrumentation.llamaindex", "LlamaIndexInstrumentor")

def instrument_crewai() -> bool:
    return _instrument("opentelemetry.instrumentation.crewai", "CrewAIInstrumentor")

def instrument_agno() -> bool:
    return _instrument("opentelemetry.instrumentation.agno", "AgnoInstrumentor")

# --- Custom Implementation ---

def instrument_haystack() -> bool:
    """Haystack auto-detects a configured TracerProvider at import time. We just need to ensure content tracing is enabled via env var (in impact.init())"""
    try:
        import haystack  # type: ignore[import] # noqa: F401
        logger.debug("Haystack detected — will auto-use global TracerProvider")
        return True
    except ImportError:
        logger.debug("Skipping Haystack: haystack-ai is not installed")
        return False

def instrument_pydantic_ai() -> bool:
    """Instrument Pydantic AI by calling Agent.instrument_all()."""
    try:
        from pydantic_ai import Agent  # type: ignore[import]
        Agent.instrument_all()
        logger.debug("Pydantic AI: called Agent.instrument_all() successfully")
        return True
    except ImportError:
        logger.debug("Skipping Pydantic AI: pydantic-ai is not installed")
        return False
    except Exception:
        logger.exception("Pydantic AI instrumentation failed")
        return False

# --- Infrastructure ---

def instrument_httpx() -> bool:
    """Instrument HTTPX HTTP client library for outbound HTTP request tracing."""
    return _instrument("opentelemetry.instrumentation.httpx", "HTTPXClientInstrumentor")

def instrument_azure_aisearch() -> bool:
    """Instrument Azure AI Search SDK for search query tracing."""
    return _instrument("opentelemetry.instrumentation.azure_aisearch", "AzureSearchInstrumentor")

def instrument_mcp() -> bool:
    """Instrument MCP (Model Context Protocol) for tool/resource tracing."""
    return _instrument("opentelemetry.instrumentation.mcp", "McpInstrumentor")


# --- Deprecating ---
def instrument_vertexai() -> bool:
    return _instrument("opentelemetry.instrumentation.vertexai", "VertexAIInstrumentor")


# Auto-instrument all available GenAI providers
def auto_instrument_all() -> list[bool]:
    """Auto-instrument all available GenAI providers, agent frameworks, and infrastructure."""
    results: list[bool] = []
    # --- Open Telemetry (structured span attrs) ---
    results.append(instrument_google_genai())
    results.append(instrument_openai_agents())
    # --- Self-instrumented (Microsoft) ---
    results.append(instrument_foundry_agents())
    results.append(instrument_agent_framework())
    # --- OpenLLMetry (flat span attrs — normalised by impact-metric) ---
    results.append(instrument_openai())
    results.append(instrument_aws_bedrock())
    results.append(instrument_anthropic())
    results.append(instrument_cohere())
    results.append(instrument_mistral())
    results.append(instrument_watsonx())
    results.append(instrument_groq())
    results.append(instrument_together())
    results.append(instrument_ollama())
    results.append(instrument_langchain())
    results.append(instrument_llama_index())
    results.append(instrument_crewai())
    results.append(instrument_agno())
    # --- Custom implementation ---
    results.append(instrument_haystack())
    results.append(instrument_pydantic_ai())
    # --- Infrastructure ---
    results.append(instrument_azure_aisearch())
    results.append(instrument_mcp())
    # --- Deprecating ---
    results.append(instrument_vertexai())
    return results

# Export all instrumentors
__all__ = [
    # Open Telemetry (structured span attrs)
    "instrument_openai_agents",
    "instrument_google_genai",
    # Self-instrumented (Microsoft)
    "instrument_foundry_agents",
    "instrument_agent_framework",
    # OpenLLMetry (flat span attrs)
    "instrument_openai",
    "instrument_aws_bedrock",
    "instrument_anthropic",
    "instrument_cohere",
    "instrument_mistral",
    "instrument_watsonx",
    "instrument_groq",
    "instrument_together",
    "instrument_ollama",
    "instrument_langchain",
    "instrument_llama_index",
    "instrument_crewai",
    "instrument_agno",
    # Custom implementation
    "instrument_haystack",
    "instrument_pydantic_ai",
    # Infrastructure
    "instrument_httpx",
    "instrument_azure_aisearch",
    "instrument_mcp",
    # Deprecating
    "instrument_vertexai",
    # All 
    "auto_instrument_all",
]
