"""Auto-instrumentation for popular ML frameworks.

This module provides automatic instrumentation (monkey-patching) for popular
ML frameworks like scikit-learn, XGBoost, and LightGBM, enabling zero-code
observability for model predictions.
"""

import logging
import sys
from typing import List, Optional

logger = logging.getLogger(__name__)


def autolog(
    frameworks: Optional[List[str]] = None,
    exclude: Optional[List[str]] = None,
    capture_input: bool = False,
    capture_output: bool = False,
    max_payload_size: int = 10000,
) -> None:
    """Enable automatic instrumentation for ML frameworks.

    Automatically patches popular ML frameworks to emit OpenTelemetry spans
    and metrics when making predictions. This provides zero-code instrumentation
    for models using standard framework APIs.

    IMPORTANT: MCAClient must be initialized BEFORE calling autolog().
    The autolog() function uses MCAClient.get_instance() to access the active client.

    Args:
        frameworks: List of specific frameworks to instrument (e.g., ["sklearn", "xgboost"]).
            If None, automatically detects and instruments all available frameworks.
        exclude: List of frameworks to exclude from instrumentation (e.g., ["lightgbm"]).
        capture_input: Whether to capture input data in spans (default: False for security).
        capture_output: Whether to capture output data in spans (default: False for security).
        max_payload_size: Maximum size in bytes for captured input/output data (default: 10000).
            Prevents memory bloat from large batch predictions.

    Raises:
        RuntimeError: If MCAClient has not been initialized before calling autolog().
        ValueError: If invalid framework names provided in frameworks or exclude parameters.

    Examples:
        Enable autolog for all available frameworks:
        >>> from mca_sdk import MCAClient, autolog
        >>> client = MCAClient(service_name="ml-model", model_id="mdl-001", team_name="ds")
        >>> autolog()  # Automatically instruments sklearn, xgboost, lightgbm if installed

        Enable autolog for specific frameworks only:
        >>> autolog(frameworks=["sklearn", "xgboost"])

        Exclude specific frameworks:
        >>> autolog(exclude=["lightgbm"])

        Capture input/output data with size limits (use with caution - PHI risks):
        >>> autolog(capture_input=True, capture_output=True, max_payload_size=5000)

    Supported Frameworks:
        - sklearn: scikit-learn estimators (predict, predict_proba, fit)
        - xgboost: XGBoost Booster and sklearn API (predict)
        - lightgbm: LightGBM Booster (predict)

    Note:
        TensorFlow and PyTorch are not currently supported but are planned for future releases.
    """
    from ..core.client import MCAClient

    # Validate framework names
    SUPPORTED_FRAMEWORKS = {"sklearn", "xgboost", "lightgbm"}

    if frameworks is not None:
        invalid_frameworks = set(frameworks) - SUPPORTED_FRAMEWORKS
        if invalid_frameworks:
            raise ValueError(
                f"Invalid framework names: {invalid_frameworks}. "
                f"Supported frameworks: {SUPPORTED_FRAMEWORKS}"
            )

    if exclude is not None:
        invalid_exclude = set(exclude) - SUPPORTED_FRAMEWORKS
        if invalid_exclude:
            raise ValueError(
                f"Invalid framework names in exclude: {invalid_exclude}. "
                f"Supported frameworks: {SUPPORTED_FRAMEWORKS}"
            )

    # Check that MCAClient has been initialized
    client_instance = MCAClient.get_instance()
    if client_instance is None:
        raise RuntimeError(
            "MCAClient must be initialized before calling autolog(). "
            "Example: client = MCAClient(service_name='...', model_id='...', team_name='...')"
        )

    # Determine which frameworks to patch
    if frameworks is None:
        # Auto-detect all available frameworks
        frameworks_to_patch = _detect_installed_frameworks()
    else:
        # Use explicitly requested frameworks
        frameworks_to_patch = frameworks

    # Remove excluded frameworks
    if exclude:
        frameworks_to_patch = [f for f in frameworks_to_patch if f not in exclude]

    if not frameworks_to_patch:
        logger.warning("No ML frameworks detected or all frameworks excluded. Autolog is a no-op.")
        return

    logger.info(f"Enabling autolog for frameworks: {frameworks_to_patch}")

    # Patch each framework
    for framework_name in frameworks_to_patch:
        try:
            _patch_framework(framework_name, capture_input, capture_output, max_payload_size)
            logger.info(f"Successfully patched {framework_name}")
        except Exception as e:
            # Graceful failure: log warning but continue
            logger.warning(
                f"Failed to patch {framework_name}: {e}. Continuing without instrumentation for this framework."
            )


def _detect_installed_frameworks() -> List[str]:
    """Detect which ML frameworks are currently imported.

    Returns:
        List of framework names that are available in sys.modules.
    """
    available = []

    # Check for sklearn (may be imported as 'sklearn' or 'sklearn.base')
    if "sklearn" in sys.modules or "sklearn.base" in sys.modules:
        available.append("sklearn")

    # Check for xgboost
    if "xgboost" in sys.modules:
        available.append("xgboost")

    # Check for lightgbm
    if "lightgbm" in sys.modules:
        available.append("lightgbm")

    return available


def _patch_framework(
    framework_name: str,
    capture_input: bool,
    capture_output: bool,
    max_payload_size: int,
) -> None:
    """Patch a specific ML framework.

    Args:
        framework_name: Name of the framework to patch ("sklearn", "xgboost", "lightgbm").
        capture_input: Whether to capture input data.
        capture_output: Whether to capture output data.
        max_payload_size: Maximum size in bytes for captured data.

    Raises:
        ValueError: If framework_name is not recognized.
        ImportError: If framework module cannot be imported.
    """
    if framework_name == "sklearn":
        from ._sklearn import patch_sklearn

        patch_sklearn(capture_input, capture_output, max_payload_size)
    elif framework_name == "xgboost":
        from ._xgboost import patch_xgboost

        patch_xgboost(capture_input, capture_output, max_payload_size)
    elif framework_name == "lightgbm":
        from ._lightgbm import patch_lightgbm

        patch_lightgbm(capture_input, capture_output, max_payload_size)
    else:
        raise ValueError(
            f"Unknown framework: {framework_name}. Supported: sklearn, xgboost, lightgbm"
        )


__all__ = ["autolog"]
