import sys
from pathlib import Path

import numpy as np
import pandas as pd
import xgboost
from sklearn.preprocessing import LabelEncoder
from ugbio_core.logger import logger


def load_xgb_model(xgb_model_path: Path) -> "xgboost.XGBClassifier":
    """
    Load a pre-trained XGBoost model from a path.

    Parameters
    ----------
    xgb_model_path : Path
        Path to the XGBoost model.

    Returns
    -------
    xgboost.XGBClassifier
        The loaded XGBoost classifier model.
    """
    # load xgb model
    xgb_clf_es = xgboost.XGBClassifier()
    xgb_clf_es.load_model(xgb_model_path)
    return xgb_clf_es


def set_categorical_columns(df):
    """
    Convert categorical columns in DataFrame to encoded categories.

    This function identifies categorical columns (object and category dtypes),
    applies label encoding to convert them to numerical values, and ensures
    all object columns are converted to category dtype.

    Parameters
    ----------
    df : pd.DataFrame
        Input DataFrame to process. Modified in place.

    Notes
    -----
    This function modifies the DataFrame in place. Categorical columns are
    label-encoded using sklearn.preprocessing.LabelEncoder.
    """
    categorical_columns = df.select_dtypes(include=["object", "category"]).columns
    le = LabelEncoder()
    for col in categorical_columns:
        df.loc[:, col] = le.fit_transform(df[col].astype(str))
    for col in df.select_dtypes(include="object").columns:
        df[col] = df[col].astype("category")


def predict(xgb_model: "xgboost.XGBClassifier", df_calls: "pd.DataFrame") -> "np.ndarray":
    """
    Generate prediction probabilities for the positive class using a trained XGBoost model.

    Parameters
    ----------
    xgb_model : xgboost.XGBClassifier
        Trained XGBoost classifier with accessible feature names.
    df_calls : pd.DataFrame
        DataFrame containing feature columns required by the model.

    Returns
    -------
    np.ndarray
        Array of predicted probabilities for the positive class ("1").

    Raises
    ------
    SystemExit
        If required model features are missing from the input DataFrame.
    """
    model_features = xgb_model.get_booster().feature_names
    missing_features = [f for f in model_features if f not in df_calls.columns]
    if len(missing_features) > 0:
        logger.error(f"The following model features are missing from input DataFrame: {missing_features}")
        sys.exit(1)
    X = df_calls[model_features]  # noqa: N806

    set_categorical_columns(X)

    return xgb_model.predict_proba(X)[:, 1]
