"""LightGBM auto-instrumentation.

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

import logging

from ._base import BasePatcher

logger = logging.getLogger(__name__)


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

    Instruments the following LightGBM APIs:
    - lightgbm.Booster.predict(): Native API predictions
    - lightgbm.sklearn.LGBMClassifier.predict(): Sklearn-compatible classifier
    - lightgbm.sklearn.LGBMRegressor.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 lightgbm is not installed.
    """
    try:
        import lightgbm  # noqa: F401
    except ImportError:
        raise ImportError("lightgbm is not installed. Install it with: pip install lightgbm")

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


class LightGBMPatcher(BasePatcher):
    """Patcher for LightGBM framework."""

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

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

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

        # Patch sklearn-compatible API (if available)
        try:
            from lightgbm.sklearn import LGBMClassifier

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

        try:
            from lightgbm.sklearn import LGBMRegressor

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

        # Patch predict_proba for classifiers
        try:
            from lightgbm.sklearn import LGBMClassifier

            self.patch_method(
                target_class=LGBMClassifier,
                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("lightgbm.sklearn.LGBMClassifier.predict_proba not found, skipping")

        logger.info("Successfully patched lightgbm models")
