from flask import Flask, request, jsonify
import pandas as pd
import joblib
import numpy as np

# Path to the serialized model (same as in the notebook)
MODEL_PATH = "superkart_best_model.joblib"

# Load model at startup
model = joblib.load(MODEL_PATH)

# Feature names used during training
try:
    FEATURE_NAMES = list(model.feature_names_in_)
except AttributeError:
    # Fallback: user must ensure the incoming data has correct columns
    FEATURE_NAMES = None

app = Flask(__name__)

@app.route("/", methods=["GET"])
def home():
    return jsonify({"message": "SuperKart Sales Forecasting API is running."})

@app.route("/predict", methods=["POST"])
def predict():
    """
    Expected JSON format:
    {
        "data": {
            "Product_Weight": 10.5,
            "Product_Sugar_Content": "Low",
            ...
        }
    }
    or
    {
        "data": [
            {...},
            {...}
        ]
    }
    """
    payload = request.get_json()

    if payload is None or "data" not in payload:
        return jsonify({"error": "Request JSON must contain a 'data' field."}), 400

    data = payload["data"]
    if isinstance(data, dict):
        data = [data]

    df_input = pd.DataFrame(data)

    # Ensure column order matches training
    if FEATURE_NAMES is not None:
        df_input = df_input.reindex(columns=FEATURE_NAMES)

    preds = model.predict(df_input)
    preds = preds.tolist()

    return jsonify({"predictions": preds})

if __name__ == "__main__":
    # Run on port 7860 for Docker/HF
    app.run(host="0.0.0.0", port=7860)