"""
Import hooks for automatic instrumentation.

This module provides automatic instrumentation of LLM provider SDKs
by wrapping builtins.__import__ to intercept imports and apply
instrumentation after modules are loaded.

Based on ARCHITECTURE_V2.md Section 6: Auto-Instrumentation Engine.

Usage:
    # Automatic via environment variable:
    RISICARE_TRACING=true python app.py

    # Programmatic:
    from risicare.instrumentation.hooks import install_import_hooks
    install_import_hooks()

    import openai  # Automatically instrumented
"""

from __future__ import annotations

import builtins
import importlib
import logging
import sys
import warnings
from typing import Any, Callable, Dict, Optional, Set

logger = logging.getLogger(__name__)


# Type alias for instrumentor functions
Instrumentor = Callable[[object], None]

# Module name -> instrumentation function mapping
# Functions are lazy-loaded from providers submodule
_INSTRUMENTORS: Optional[Dict[str, str]] = None


def _get_instrumentors() -> Dict[str, str]:
    """
    Get the instrumentor function mapping.

    Maps module names to the function path that instruments them.
    Functions are lazy-loaded to avoid import cycles.
    """
    global _INSTRUMENTORS
    if _INSTRUMENTORS is None:
        _INSTRUMENTORS = {
            # LLM Providers (original 5)
            "openai": "risicare.instrumentation.providers.openai:instrument_openai",
            "anthropic": "risicare.instrumentation.providers.anthropic:instrument_anthropic",
            "cohere": "risicare.instrumentation.providers.cohere:instrument_cohere",
            "google.generativeai": "risicare.instrumentation.providers.google:instrument_google",
            "mistralai": "risicare.instrumentation.providers.mistral:instrument_mistral",
            # LLM Providers (GAP-07 additions)
            "groq": "risicare.instrumentation.providers.groq:instrument_groq",
            "together": "risicare.instrumentation.providers.together:instrument_together",
            "ollama": "risicare.instrumentation.providers.ollama:instrument_ollama",
            "botocore": "risicare.instrumentation.providers.bedrock:instrument_bedrock",
            "vertexai": "risicare.instrumentation.providers.vertex_ai:instrument_vertex_ai",
            # LLM Providers (extended provider parity)
            "cerebras": "risicare.instrumentation.providers.cerebras:instrument_cerebras",
            "huggingface_hub": "risicare.instrumentation.providers.huggingface:instrument_huggingface",
            # Agent Frameworks (in-SDK integrations)
            "langchain": "risicare.integrations.langchain:instrument_langchain",
            "langchain_core": "risicare.integrations.langchain:instrument_langchain",
            "langgraph": "risicare.integrations.langgraph:instrument_langgraph",
            "crewai": "risicare.integrations.crewai:instrument_crewai",
            "autogen": "risicare.integrations.autogen:instrument_autogen",
            "autogen_agentchat": "risicare.integrations.autogen:instrument_autogen",
            "agents": "risicare.integrations.openai_agents:instrument_openai_agents",
            # AI Frameworks (Tier 1 additions)
            "instructor": "risicare.integrations.instructor:instrument_instructor",
            "litellm": "risicare.integrations.litellm:instrument_litellm",
            "dspy": "risicare.integrations.dspy:instrument_dspy",
            "pydantic_ai": "risicare.integrations.pydantic_ai:instrument_pydantic_ai",
            "llama_index": "risicare.integrations.llamaindex:instrument_llamaindex",
            "llama_index.core": "risicare.integrations.llamaindex:instrument_llamaindex",
            # OpenTelemetry bridge (opt-in via otel_bridge=True)
            "opentelemetry": "risicare.integrations.otel:instrument_opentelemetry",
        }
    return _INSTRUMENTORS


def _load_instrumentor(path: str) -> Optional[Instrumentor]:
    """
    Load an instrumentor function from a module path.

    Args:
        path: Path in format "module.path:function_name"

    Returns:
        The instrumentor function or None if loading fails.
    """
    try:
        module_path, func_name = path.rsplit(":", 1)
        module = importlib.import_module(module_path)
        return getattr(module, func_name)
    except Exception as e:
        logger.debug(f"Failed to load instrumentor {path}: {e}")
        return None


# Kept for backward compatibility (referenced in __init__.py exports)
class RisicareImportFinder:
    """Tracks instrumented modules. No longer used as a MetaPathFinder."""

    def __init__(self, allowed_modules: Optional[Set[str]] = None) -> None:
        instrumentors = _get_instrumentors()
        self.allowed_modules = allowed_modules or set(instrumentors.keys())
        self._instrumented: Set[str] = set()


# Global state
_finder: Optional[RisicareImportFinder] = None
_original_import: Optional[Any] = None
_hooks_installed: bool = False


def _instrumented_import(name: str, *args: Any, **kwargs: Any) -> Any:
    """
    Replacement for builtins.__import__ that instruments LLM provider modules
    after they are loaded by the normal import system.
    """
    result = _original_import(name, *args, **kwargs)

    # Check if this is a module we should instrument
    instrumentors = _get_instrumentors()
    base_module = name.split(".")[0]

    # Check if base module or full name matches
    target = None
    if name in instrumentors:
        target = name
    elif base_module in instrumentors:
        target = base_module

    if target and _finder and target not in _finder._instrumented:
        # Module is now in sys.modules, instrument it
        if target in sys.modules:
            module = sys.modules[target]
            instrumentor_path = instrumentors[target]
            instrumentor = _load_instrumentor(instrumentor_path)
            if instrumentor:
                try:
                    instrumentor(module)
                    _finder._instrumented.add(target)
                    logger.debug(f"Instrumented {target}")
                except Exception as e:
                    warnings.warn(
                        f"Failed to instrument {target}: {e}",
                        RuntimeWarning,
                    )

    return result


def install_import_hooks(
    allowed_modules: Optional[Set[str]] = None,
) -> RisicareImportFinder:
    """
    Install instrumentation import hooks.

    Wraps builtins.__import__ to automatically instrument LLM provider
    modules when they are imported.

    Args:
        allowed_modules: Set of module names to instrument.
                        If None, all supported modules are instrumented.

    Returns:
        The tracker instance.
    """
    global _finder, _original_import, _hooks_installed

    if _hooks_installed:
        if allowed_modules and _finder and allowed_modules != _finder.allowed_modules:
            logger.warning(
                "import_hooks already installed with different allowed_modules. "
                "Call remove_import_hooks() first to change modules."
            )
        if _finder is None:
            _finder = RisicareImportFinder(allowed_modules)
        return _finder

    _finder = RisicareImportFinder(allowed_modules)

    # Store original import ONCE — never overwrite on reinstall
    if _original_import is None:
        _original_import = builtins.__import__
    builtins.__import__ = _instrumented_import
    _hooks_installed = True
    logger.debug("Installed Risicare import hooks")

    # Instrument any already-imported modules
    instrument_already_imported()

    return _finder


def remove_import_hooks() -> None:
    """Remove instrumentation import hooks."""
    global _finder, _original_import, _hooks_installed

    if _original_import is not None:
        builtins.__import__ = _original_import
        _original_import = None

    _finder = None
    _hooks_installed = False
    logger.debug("Removed Risicare import hooks")


def instrument_already_imported() -> int:
    """
    Instrument modules that are already imported.

    This handles the case where modules are imported before hooks are installed.

    Returns:
        Number of modules instrumented.
    """
    instrumented_count = 0
    instrumentors = _get_instrumentors()

    for module_name, instrumentor_path in instrumentors.items():
        if module_name in sys.modules:
            module = sys.modules[module_name]

            # Skip if already instrumented
            if _finder and module_name in _finder._instrumented:
                continue

            instrumentor = _load_instrumentor(instrumentor_path)
            if instrumentor:
                try:
                    instrumentor(module)
                    if _finder:
                        _finder._instrumented.add(module_name)
                    instrumented_count += 1
                    logger.debug(f"Instrumented already-imported {module_name}")
                except Exception as e:
                    warnings.warn(
                        f"Failed to instrument already-imported {module_name}: {e}",
                        RuntimeWarning,
                    )

    return instrumented_count


def is_instrumented(module_name: str) -> bool:
    """Check if a module has been instrumented."""
    if _finder is None:
        return False
    return module_name in _finder._instrumented


def get_instrumented_modules() -> Set[str]:
    """Get the set of instrumented module names."""
    if _finder is None:
        return set()
    return _finder._instrumented.copy()


def get_supported_modules() -> Set[str]:
    """Get the set of module names that can be instrumented."""
    return set(_get_instrumentors().keys())
