import json
import joblib
import numpy as np
from lightgbm import LGBMRegressor

# Load the model once when the endpoint starts
model = joblib.load("lgbm_model.joblib")  # <-- Your saved model

def preprocess(inputs):
    """
    Convert JSON input into the appropriate format for LightGBM model
    inputs: dict or list of dicts
    """
    # If input is a single dict, convert it to a list
    if isinstance(inputs, dict):
        inputs = [inputs]
    
    # Convert to numpy array (feature order must match training!)
    feature_order = sorted(inputs[0].keys())
    X = np.array([[sample[f] for f in feature_order] for sample in inputs])
    return X

def predict(inputs):
    X = preprocess(inputs)
    preds = model.predict(X)
    # Return a list of predictions
    return preds.tolist()

def handle(inputs):
    """
    Function called by Hugging Face Inference Endpoint
    """
    # If input is a JSON string
    if isinstance(inputs, str):
        inputs = json.loads(inputs)
    
    preds = predict(inputs)
    return {"predictions": preds}