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

# Initialize Flask app
app = Flask(__name__)

# Load the trained sales forecasting model pipeline
model = joblib.load("backend_files/final_sales_forecasting_model.joblib") # Adjust path as needed

# Define a route for the home page
@app.route('/')
def home():
    return "Welcome to the SuperKart Sales Forecasting API"

# Define an endpoint to predict sales for a single product-store combination
@app.route('/predict_single', methods=['POST'])
def predict_single():
    # Get JSON data from the request
    data = request.get_json()

    # Extract relevant features from the input data, ensuring correct order and names
    # The keys in the dictionary should match the column names in your original training data X
    try:
        sample = {
            'Product_Id': data['Product_Id'],
            '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_Establishment_Year': data['Store_Establishment_Year'],
            'Store_Size': data['Store_Size'],
            'Store_Location_City_Type': data['Store_Location_City_Type'],
            'Store_Type': data['Store_Type']
        }

        # Convert the extracted data into a DataFrame
        input_data = pd.DataFrame([sample])

        # Make a sales prediction using the trained model pipeline
        prediction = model.predict(input_data).tolist()[0]

        # Return the prediction as a JSON response
        return jsonify({'predicted_sales': prediction})

    except KeyError as e:
        return jsonify({'error': f'Missing data for key: {e}'}), 400
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# Define an endpoint to predict sales for a batch of product-store combinations from a CSV file
@app.route('/predict_batch', methods=['POST'])
def predict_batch():
    # Get the uploaded file from the request
    if 'file' not in request.files:
        return jsonify({'error': 'No file part in the request'}), 400

    file = request.files['file']

    # If the user does not select a file, the browser submits an empty file without a filename.
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    if file:
        try:
            # Read the file into a DataFrame
            input_data = pd.read_csv(file)

            # Make sales predictions using the trained model pipeline
            predictions = model.predict(input_data).tolist()

            # Return the predictions as a JSON response
            return jsonify({'predicted_sales': predictions})

        except Exception as e:
            return jsonify({'error': str(e)}), 500
    else:
        return jsonify({'error': 'Something went wrong with file upload'}), 500


if __name__ == '__main__':
    # Run the Flask app
    app.run(debug=True, host='0.0.0.0', port=5000) # Run on all available interfaces and port 5000