from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import pipeline
import os

print("Loading End-to-End question generation model...")
question_generator = pipeline("text2text-generation", model="valhalla/t5-base-e2e-qg")
print("Model loaded successfully!")

app = Flask(__name__)
CORS(app)

@app.route('/generate-quiz', methods=['POST'])
def generate_quiz():
    data = request.get_json()
    if not data or 'context' not in data:
        return jsonify({'error': 'No context provided in the request.'}), 400

    context = data['context']
    print(f"Generating questions for the provided context...")

    generated_questions = []
    try:
        output = question_generator(
            context,
            max_length=64,
            num_beams=5,
            num_return_sequences=5,
            early_stopping=True
        )

        question_set = set()
        for item in output:
            question_str = item['generated_text']
            questions = question_str.split('<sep>')
            for q in questions:
                clean_q = q.strip()
                if clean_q.endswith('?') and len(clean_q.split()) > 3:
                    question_set.add(clean_q)

        generated_questions = [{'question': q, 'context': context} for q in question_set]
        generated_questions = generated_questions[:10]

    except Exception as e:
        print(f"Error generating questions: {e}")
        return jsonify({'error': 'Failed to generate quiz.'}), 500

    print(f"Successfully generated {len(generated_questions)} unique questions.")
    return jsonify({'quiz': generated_questions})

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 5000))
    app.run(host='0.0.0.0', port=port)