import gradio as gr
import google.generativeai as genai
import os
import time

genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
gemini_model = genai.GenerativeModel("gemini-2.0-flash")

def prepare_ranking_prompt(question, correct, distractors):
    distractor_text = "\n".join([f"Choice {i+1}: {d}" for i, d in enumerate(distractors)])
    return f"""You are a helpful assistant evaluating multiple-choice distractors.

Question: {question}
Correct Answer: {correct}

Here are the distractors:
{distractor_text}

Your task: Rank the distractors from most to least confusing.

Only respond with one line in this format: 2 > 1 > 3
"""

def generate_and_rank(question):
    gemini_prompt = f"""You are an assistant generating multiple-choice questions.

Given a question, generate:
- One correct answer
- Three plausible but incorrect distractors

Format:
Correct Answer: ...
Distractor 1: ...
Distractor 2: ...
Distractor 3: ...

Question: {question}
""" 
    gemini_response = gemini_model.generate_content(gemini_prompt).text.strip()

    # Parse output
    lines = [line.strip() for line in gemini_response.splitlines() if line.strip()]
    correct = next((line.split(":", 1)[1].strip() for line in lines if line.startswith("Correct Answer")), "N/A")
    distractors = [line.split(":", 1)[1].strip() for line in lines if line.startswith("Distractor")]

    if len(distractors) != 3:
        return gemini_response, "ERROR (Invalid distractor count)"

    ranking_prompt = prepare_ranking_prompt(question, correct, distractors)
    ranking_response = gemini_model.generate_content(ranking_prompt).text.strip()

    return gemini_response, ranking_response

iface = gr.Interface(
    fn=generate_and_rank,
    inputs=gr.Textbox(label="Enter your MCQ Question"),
    outputs=[
        gr.Textbox(label="Generated Answer and Distractors"),
        gr.Textbox(label="Ranking of Distractors")
    ],
    title="Confusing Distractor Generator + Ranking"
)

iface.launch()