import joblib
import pandas as pd

class EndpointHandler:
    def __init__(self, path=""):
        # Load model dari folder repo
        self.bundle = joblib.load(f"{path}/nutrisnaps_rf_bundle.joblib")
        self.model = self.bundle['model']
        self.feature_names = self.bundle['feature_names']
        self.impute_col = self.bundle['impute_col']

    def __call__(self, data):
        # Ambil input JSON
        inputs = data.pop("inputs", data)
        df_input = pd.DataFrame([inputs])
        
        # Penyesuaian kolom fitur
        for col in self.feature_names:
            if col not in df_input.columns:
                df_input[col] = 0 
                
        df_input = df_input[self.feature_names]

        # Prediksi
        prediction = self.model.predict(df_input)
        return {"prediction": prediction.tolist()}