"""Scikit-learn auto-instrumentation.

This module provides automatic instrumentation for scikit-learn models
by patching the BaseEstimator class methods.
"""

import logging

from ._base import BasePatcher

logger = logging.getLogger(__name__)


def patch_sklearn(
    capture_input: bool = False,
    capture_output: bool = False,
    max_payload_size: int = 10000,
) -> None:
    """Patch scikit-learn estimators for automatic instrumentation.

    Instruments the following methods on sklearn.base.BaseEstimator:
    - predict(): Main prediction method
    - predict_proba(): Probability predictions
    - fit(): Model training (optional, for training time tracking)

    Args:
        capture_input: Whether to capture input data shapes in spans.
        capture_output: Whether to capture output data shapes in spans.
        max_payload_size: Maximum size in bytes for captured data.

    Raises:
        ImportError: If scikit-learn is not installed.
    """
    try:
        from sklearn.base import BaseEstimator  # noqa: F401
    except ImportError:
        raise ImportError(
            "scikit-learn is not installed. Install it with: pip install scikit-learn"
        )

    patcher = SklearnPatcher(capture_input, capture_output, max_payload_size)
    patcher.patch()


class SklearnPatcher(BasePatcher):
    """Patcher for scikit-learn framework."""

    def __init__(
        self,
        capture_input: bool = False,
        capture_output: bool = False,
        max_payload_size: int = 10000,
    ):
        super().__init__("sklearn", capture_input, capture_output, max_payload_size)

    def patch(self) -> None:
        """Apply patches to sklearn.base.BaseEstimator."""
        from sklearn.base import BaseEstimator  # noqa: F401

        # Patch predict() method (most common)
        try:
            self.patch_method(
                target_class=BaseEstimator,
                method_name="predict",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={"prediction_type": "class"},
            )
        except AttributeError:
            logger.warning("sklearn.base.BaseEstimator.predict not found, skipping")

        # Patch predict_proba() method (for classifiers)
        try:
            self.patch_method(
                target_class=BaseEstimator,
                method_name="predict_proba",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={"prediction_type": "probability"},
            )
        except AttributeError:
            logger.debug("sklearn.base.BaseEstimator.predict_proba not found, skipping")

        # Patch fit() method (for training tracking)
        try:
            self.patch_method(
                target_class=BaseEstimator,
                method_name="fit",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={"operation_type": "training"},
            )
        except AttributeError:
            logger.debug("sklearn.base.BaseEstimator.fit not found, skipping")

        logger.info("Successfully patched sklearn.base.BaseEstimator")
