import joblib
import pandas as pd
from flask import Flask, request, jsonify
import sys

# Re-define the same functions used inside the pipeline
def add_store_age(df):
    df = df.copy()
    df['Store_Age'] = 2025 - df['Store_Established_Year']
    df = df.drop(columns=['Store_Established_Year'])
    return df

def map_ordered_features(X):
    sugar_order_map = {'No Sugar': 0, 'Low Sugar': 1, 'Regular': 2}
    size_order_map = {'Small': 0, 'Medium': 1, 'High': 2}
    city_order_map = {'Tier 3': 0, 'Tier 2': 1, 'Tier 1': 2}

    X = X.copy()
    X['Product_Sugar_Content'] = X['Product_Sugar_Content'].map(sugar_order_map).astype(int)
    X['Store_Size'] = X['Store_Size'].map(size_order_map).astype(int)
    X['Store_Location_City_Type'] = X['Store_Location_City_Type'].map(city_order_map).astype(int)
    return X

# Inject functions into the module namespace where joblib expects them
if __name__ != '__main__':
    # When running in production (Hugging Face), inject into __main__
    import __main__
    __main__.add_store_age = add_store_age
    __main__.map_ordered_features = map_ordered_features

# Initialize Flask app with a name
sales_predictor_api = Flask("SuperKart Sales Predictor")

# Load the trained SuperKart Sales prediction model
model = joblib.load("superkart_sales_model.pkl")

# Define a route for the home page
@sales_predictor_api.get('/')
def home():
    return "Welcome to the SuperKart Sales Prediction API!"

# Define an endpoint to predict churn for a single customer
@sales_predictor_api.post('/v1/productsales')
def predict_sales():
    try:
        # Get JSON data
        sales_data = request.get_json()

        # Ensure input is a list of records
        if isinstance(sales_data, dict):
            sales_data = [sales_data]

        # Convert to DataFrame
        input_data = pd.DataFrame(sales_data)

        # Add Store_Id if not present (model expects it but doesn't use it)
        if 'Store_Id' not in input_data.columns:
            input_data['Store_Id'] = 'DUMMY_STORE'

        # Predict
        prediction = model.predict(input_data).tolist()[0]

        return jsonify({"prediction": float(prediction)})

    except Exception as e:
        return jsonify({"error": str(e)})

# Run the Flask app in debug mode
if __name__ == '__main__':
    sales_predictor_api.run(debug=True)