from flask import Flask, request, jsonify, render_template
from huggingface_hub import InferenceClient
from requests.exceptions import RequestException

app = Flask(__name__)

client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")

def system_instructions(context):
    return f"""<s> [INST] Your are a great teacher and your task is to create 6 questions with answer and 4 choices based on the following context:\n\n{context}\n\n. Each example should be like this
    Question: ""
    Choices:
    A): ""
    B): ""
    C): ""
    D): ""
    Answer: "Choose either A, B, C, or D as the correct answer"
    Explanation: ""
    \n
    [/INST]
    """

def generate_quiz(context):
    formatted_prompt = system_instructions(context)

    generate_kwargs = dict(
            temperature=0.1,
            max_new_tokens=2048,
            top_p=0.95,
            repetition_penalty=1.0,
            do_sample=True,
            seed=42,)
        
    try:
        response = client.text_generation(
            formatted_prompt,
            **generate_kwargs,
            stream=False,
            details=False,
            return_full_text=False,
        )
        return response
        
    except (RequestException, SystemExit) as e:
        return {"error": str(e)}

@app.route("/", methods=["GET", "POST"])
def generate_quiz_page():
    if request.method == "POST":
        if request.content_type == 'application/json':
            data = request.get_json()
            context = data.get("context")
            if context is None or context.strip() == "":
                return jsonify({"error": "Missing or empty 'context' parameter"}), 400
        else:
            context = request.form.get("context")
            if context is None or context.strip() == "":
                return jsonify({"error": "Missing or empty 'context' parameter"}), 400
        
        response = generate_quiz(context)

        if request.content_type == 'application/json':
            return jsonify(response)

        # Handle text processing for quiz generation
        filtered_text, answers_and_explanations = handle_text(response)

        return render_template('quiz.html', filtered_text=filtered_text, answers_and_explanations=answers_and_explanations)
    
    # Default GET request handling
    return render_template('index.html')

def handle_text(text):
    lines = text.splitlines()
    filtered_lines = []
    answers_and_explanations = []

    for line in lines:
        line = line.strip()

        if line.startswith("Question:"):
            filtered_lines.append(line)
        elif line.startswith("Choices:"):
            filtered_lines.append(line)
        elif line.startswith("Answer:"):
            answer = line.split(":")[1].strip()
            answers_and_explanations.append({"Answer": answer})
        elif line.startswith("Explanation:"):
            explanation = line.split(":")[1].strip()
            answers_and_explanations[-1]["Explanation"] = explanation
        elif line.startswith(("A)", "B)", "C)", "D)")):
            filtered_lines.append(line)
        else:
            filtered_lines.append(line)

    filtered_text = "\n".join(filtered_lines)
    return filtered_text, answers_and_explanations

if __name__ == "__main__":
    app.run(debug=True)