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

# Initialize the Flask app
superkart_api = Flask("SuperKart Sales Prediction API")

# Load the trained model (assumed to include preprocessing pipeline)
model = joblib.load("deployment_files/best_random_forest_model.joblib")

@superkart_api.get('/')
def home():
    return "Welcome to the SuperKart Sales Prediction API!"

@superkart_api.post('/v1/sales')
def predict_sales():
    """
    Handle POST request with a JSON payload for single prediction.
    """
    try:
        input_data = request.get_json()

        # Wrap in list to convert to DataFrame
        df = pd.DataFrame([input_data])

        # Predict using the trained pipeline
        prediction = model.predict(df)[0]

        return jsonify({'Predicted Quarterly Sales (in INR)': round(float(prediction), 2)})

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

@superkart_api.post('/v1/salesbatch')
def predict_sales_batch():
    """
    Handle POST request with a CSV file for batch predictions.
    """
    try:
        file = request.files['file']
        input_df = pd.read_csv(file)

        # Predict using the trained pipeline
        predictions = model.predict(input_df)
        predictions = [round(float(val), 2) for val in predictions]

        # Add to original DataFrame
        input_df['Predicted_Quarterly_Sales'] = predictions

        return input_df.to_dict(orient='records')

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

if __name__ == '__main__':
    superkart_api.run(debug=True)