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

# Initialize Flask app
sales_api = Flask("Product Store Sales Predictor")

# Load the trained sales prediction model
model = joblib.load("sales_prediction_model_v1_0.joblib")


# ==============================
# Home route
# ==============================
@sales_api.get('/')
def home():
    return "Welcome to the Product Store Sales Prediction API!"


# ==============================
# Single record prediction endpoint
# ==============================
@sales_api.post('/v1/sales')
def predict_sales():
    try:
        # Get JSON data from the request
        data = request.get_json()

        # Extract all required product/store features from the input data
        sample = {
            'Product_Code': data['Product_Code'],
            'Product_Weight': data['Product_Weight'],
            'Product_Sugar_Content': data['Product_Sugar_Content'],
            'Product_Allocated_Area': data['Product_Allocated_Area'],
            'Product_Type': data['Product_Type'],
            'Product_MRP': data['Product_MRP'],
            'Store_Id': data['Store_Id'],
            'Store_Age': data['Store_Age'],
            'Store_Size': data['Store_Size'],
            'Store_Location_City_Type': data['Store_Location_City_Type'],
            'Store_Type': data['Store_Type']
        }

        # Convert dictionary to DataFrame
        input_df = pd.DataFrame([sample])

        # Make prediction using the trained model
        prediction = model.predict(input_df).tolist()[0]

        # Return JSON response
        return jsonify({'Predicted_Product_Store_Sales_Total': prediction})

    except Exception as e:
        return jsonify({'error': str(e)}), 400


# ==============================
# Batch prediction endpoint
# ==============================
@sales_api.post('/v1/salesbatch')
def predict_sales_batch():
    try:
        # Get uploaded CSV file from the request
        file = request.files['file']

        # Read CSV file into DataFrame
        input_data = pd.read_csv(file)

        # Make sure all required columns exist
        required_cols = [
            'Product_Code', 'Product_Weight',
            'Product_Sugar_Content', 'Product_Allocated_Area', 'Product_Type',
            'Product_MRP', 'Store_Id', 'Store_Age',
            'Store_Size', 'Store_Location_City_Type', 'Store_Type'
        ]

        missing_cols = set(required_cols) - set(input_data.columns)
        if missing_cols:
            return jsonify({'error': f"columns are missing: {missing_cols}"}), 400

        # Make predictions on the batch
        predictions = model.predict(input_data).tolist()

        # Add predictions as new column
        input_data['Predicted_Product_Store_Sales_Total'] = predictions

        # Convert DataFrame to list of dicts for JSON output
        result = input_data.to_dict(orient='records')

        return jsonify(result)

    except Exception as e:
        return jsonify({'error': str(e)}), 400


# ==============================
# Run Flask app
# ==============================
if __name__ == '__main__':
    sales_api.run(debug=True)