"""XGBoost auto-instrumentation.

This module provides automatic instrumentation for XGBoost models
by patching both the native Booster API and sklearn-compatible API.
"""

import logging

from ._base import BasePatcher

logger = logging.getLogger(__name__)


def patch_xgboost(
    capture_input: bool = False,
    capture_output: bool = False,
    max_payload_size: int = 10000,
) -> None:
    """Patch XGBoost models for automatic instrumentation.

    Instruments the following XGBoost APIs:
    - xgboost.Booster.predict(): Native API predictions
    - xgboost.sklearn.XGBClassifier.predict(): Sklearn-compatible classifier
    - xgboost.sklearn.XGBRegressor.predict(): Sklearn-compatible regressor

    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 xgboost is not installed.
    """
    try:
        import xgboost  # noqa: F401
    except ImportError:
        raise ImportError("xgboost is not installed. Install it with: pip install xgboost")

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


class XGBoostPatcher(BasePatcher):
    """Patcher for XGBoost framework."""

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

    def patch(self) -> None:
        """Apply patches to XGBoost classes."""
        import xgboost  # noqa: F401

        # Patch native Booster API
        try:
            self.patch_method(
                target_class=xgboost.Booster,
                method_name="predict",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={"api_type": "native"},
            )
        except AttributeError:
            logger.warning("xgboost.Booster.predict not found, skipping")

        # Patch sklearn-compatible API (if available)
        try:
            from xgboost.sklearn import XGBClassifier

            self.patch_method(
                target_class=XGBClassifier,
                method_name="predict",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={"api_type": "sklearn", "model_type": "classifier"},
            )
        except (ImportError, AttributeError):
            logger.debug("xgboost.sklearn.XGBClassifier not found, skipping")

        try:
            from xgboost.sklearn import XGBRegressor

            self.patch_method(
                target_class=XGBRegressor,
                method_name="predict",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={"api_type": "sklearn", "model_type": "regressor"},
            )
        except (ImportError, AttributeError):
            logger.debug("xgboost.sklearn.XGBRegressor not found, skipping")

        # Patch predict_proba for classifiers
        try:
            from xgboost.sklearn import XGBClassifier

            self.patch_method(
                target_class=XGBClassifier,
                method_name="predict_proba",
                span_name_template="{framework}.{class_name}.{method}",
                additional_attrs={
                    "api_type": "sklearn",
                    "model_type": "classifier",
                    "prediction_type": "probability",
                },
            )
        except (ImportError, AttributeError):
            logger.debug("xgboost.sklearn.XGBClassifier.predict_proba not found, skipping")

        logger.info("Successfully patched xgboost models")
