from flask import Flask, request, jsonify
from transformers import pipeline
import logging

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

app = Flask(__name__)

# Load the sentiment analysis pipeline from Hugging Face
# This will download the model if not already cached
logging.info("Loading sentiment analysis pipeline...")
# For the purpose of this app, we assume the model to be loaded will be 'Arvind111/bert_sentiment-analysis-model' if pushed.
# For local testing, you might use a general sentiment model, but the task specifies using the pushed model.
# Placeholder for the actual model to be loaded from HF. For now, using a general one for structure.
# In a real deployment, we'd replace 'distilbert-base-uncased-finetuned-sst-2-english' with 'Arvind111/bert_sentiment-analysis-model'
# once it's pushed and available.
classifier = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
logging.info("Sentiment analysis pipeline loaded.")

# Define the label mapping as per the trained model (0=Sad, 1=Happy, 2=Neutral)
# The simpletransformers model outputs 0, 1, 2 directly, but the transformers pipeline might output LABEL_0, LABEL_1, etc.
# We need to map these to our desired labels. The simpletransformers model's labels were 0: Sad, 1: happy, 2: Neutral
# If the pipeline returns LABEL_0, LABEL_1, LABEL_2 based on the fine-tuned model, we map accordingly.
# Assuming LABEL_0 -> Sad, LABEL_1 -> Happy, LABEL_2 -> Neutral based on the training data order.
label_mapping = {
    'LABEL_0': 'Sad',
    'LABEL_1': 'Happy',
    'LABEL_2': 'Neutral'
}

@app.route('/')
def home():
    return "Welcome to the Emotion Prediction API! Use the /predict endpoint."

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json(force=True)
    text = data.get('text', '')

    if not text:
        logging.warning("Received empty text for prediction.")
        return jsonify({'error': 'No text provided for prediction'}), 400

    try:
        # The pipeline returns a list of dictionaries, e.g., [{'label': 'LABEL_1', 'score': 0.999}]
        prediction_result = classifier(text)[0]
        predicted_label_code = prediction_result['label']
        score = prediction_result['score']

        # Map the Hugging Face label code to our custom emotion label
        emotion = label_mapping.get(predicted_label_code, 'Unknown')

        response = {
            'text': text,
            'predicted_emotion': emotion,
            'confidence': score
        }
        logging.info(f"Prediction made for text: '{text[:50]}' - Emotion: {emotion}, Confidence: {score:.4f}")
        return jsonify(response)

    except Exception as e:
        logging.error(f"Error during prediction for text: '{text[:50]}' - Error: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # Use a default port or get it from environment variables for deployment
    port = 5000
    logging.info(f"Starting Flask app on port {port}")
    app.run(host='0.0.0.0', port=port)