import os
import json
import gradio as gr
from mistralai import Mistral

api_key = os.environ["MISTRAL_API_KEY"]
client = Mistral(api_key=api_key)
model_name = "mistral-large-latest"


def generate_MCQquestion(text):
    
    prompt = f"""
    You are a strict Arabic exam generator.
    
    You MUST use ONLY the information explicitly written in the Context.
    DO NOT use any external knowledge.
    DO NOT add any information not present in the Context.
    
    Context:
    {text}
    Generate MCQs in Arabic based on the length of the Context.
    
    Question count rules:
    - Very short context (less than 100 words): generate 3 questions only
    - Short context (100–300 words): generate 3–5 questions
    - Medium context (300–700 words): generate 5–7 questions
    - Long context (more than 700 words): generate 7–10 questions
    STRICT RULES:
    - DO NOT generate more questions than the Context can support.
    - DO NOT invent information just to increase the number of questions.
    - Quality is more important than quantity.
    - Every question MUST be answerable ONLY using the Context.
    - Every correct answer MUST exist explicitly in the Context.
    - DO NOT infer or assume information.
    - DO NOT use outside knowledge.
    - DO NOT create general biology questions.
    - Use ONLY facts present in the Context text.
    - If the answer cannot be found directly in the Context, DO NOT create the question.

    
    Options rules:
    - 4 options per question
    - Only one correct answer
    - Incorrect options must be plausible but still related to the Context
    - No duplicate options
    
    Return ONLY valid JSON:
    
    {{
      "mcqs": [
        {{
          "question": "string",
          "options": ["opt1", "opt2", "opt3", "opt4"],
          "answer_index": integer
          "answer_source": "exact sentence from context"
        }}
      ]
    }}



"""

    try:

        response = client.chat.complete(
            model=model_name,
            messages=[
                {"role": "user", "content": prompt}
            ],
            response_format={"type": "json_object"},
            temperature=0.2
        )

        content = response.choices[0].message.content
        data = json.loads(content)

        # handle both dict and list formats
        if isinstance(data, dict):
            mcqs = data.get("mcqs", [])
        elif isinstance(data, list):
            mcqs = data
        else:
            return json.dumps({"error": "Invalid response format"}, ensure_ascii=False)

        # enforce limits
        if len(mcqs) < 3:
            return json.dumps({"error": "Model generated fewer than 3 questions"}, ensure_ascii=False)

        if len(mcqs) > 10:
            mcqs = mcqs[:10]

        questions = []
        answers = []

        for i, item in enumerate(mcqs, start=1):

            # safety checks
            if (
                "question" not in item or
                "options" not in item or
                "answer_index" not in item
            ):
                continue

            if not isinstance(item["options"], list) or len(item["options"]) != 4:
                continue

            if item["answer_index"] < 0 or item["answer_index"] > 3:
                continue

            questions.append({
                "id": i,
                "question": item["question"],
                "options": item["options"]
            })

            answers.append({
                "question_id": i,
                "correct_answer_index": item["answer_index"] + 1
            })

        final_output = {
            "questions": questions,
            "answers": answers
        }

        return json.dumps(final_output, ensure_ascii=False, indent=2)

    except Exception as e:
        return json.dumps({"error": str(e)}, ensure_ascii=False, indent=2)


# UI
demo = gr.Interface(
    fn=generate_MCQquestion,
    inputs=gr.Textbox(lines=10, label="أدخل النص"),
    outputs=gr.Code(language="json", label="JSON Output"),
    title="Arabic MCQ Generator API Ready",
)

demo.launch()