#!/usr/bin/env python3

import json

import os

import random

import re

import sys

import urllib.error

import urllib.request





def clean_text(value):

    return re.sub(r"\s+", " ", (value or "")).strip()





def split_topics(raw_topics):

    items = [clean_text(item) for item in re.split(r"[,;\n]+", raw_topics or "")]

    return [item for item in items if item]





def build_subject_profiles(subject_name):

    subject = (subject_name or "").lower()

    if "mathematics" in subject:

        return {

            "objective_templates": [

                "Which option correctly solves a problem on {topic}?",

                "The correct mathematical statement about {topic} is",

                "Which option shows the best reasoning for {topic}?",

            ],

            "subjective_templates": [

                "Solve a problem involving {topic} and show all workings.",

                "Explain the steps used in {topic} with a worked example.",

                "Use {topic} to answer a practical question.",

            ],

            "passage_style": "numeracy, method, and applied reasoning",

        }

    if "english" in subject:

        return {

            "objective_templates": [

                "Which option best explains the use of {topic}?",

                "Choose the sentence that correctly reflects {topic}.",

                "Which option best matches the meaning of {topic}?",

            ],

            "subjective_templates": [

                "Explain {topic} clearly and give suitable examples.",

                "Discuss {topic} using clear and correct language.",

                "Write short notes on {topic} and its importance.",

            ],

            "passage_style": "clear language, vocabulary, and understanding",

        }

    if "science" in subject:

        return {

            "objective_templates": [

                "Which option best explains {topic}?",

                "The correct scientific statement about {topic} is",

                "Which option shows the best application of {topic}?",

            ],

            "subjective_templates": [

                "Explain {topic} and relate it to everyday life.",

                "Describe {topic} and state its practical importance.",

                "Discuss the process or idea behind {topic}.",

            ],

            "passage_style": "observation, explanation, and practical application",

        }

    return {

        "objective_templates": [

            "Which of the following best explains {topic}?",

            "One correct statement about {topic} is",

            "Which option is most closely related to {topic}?",

        ],

        "subjective_templates": [

            "Explain {topic} in detail and give suitable examples.",

            "Discuss the importance of {topic} in this subject.",

            "Outline key points learners should know about {topic}.",

        ],

        "passage_style": "clear understanding and application",

    }





def extract_seed_terms(materials, retrieval_context, sample_questions, explicit_topics, subject_name):

    seeds = []

    for topic in explicit_topics:

        if topic and topic not in seeds:

            seeds.append(topic)



    for chunk in retrieval_context:

        title = clean_text(chunk.get("title"))

        if title and title not in seeds:

            seeds.append(title)

        for token in re.findall(r"\b[A-Z][a-zA-Z]{3,}(?:\s+[A-Z][a-zA-Z]{3,})?\b", chunk.get("text", "")):

            token = clean_text(token)

            if token and token.lower() != subject_name.lower() and token not in seeds:

                seeds.append(token)

            if len(seeds) >= 20:

                return seeds



    for sample in split_topics((sample_questions or "").replace("?", "\n")):

        if len(sample) > 4 and sample not in seeds:

            seeds.append(sample)



    for material in materials:

        content = clean_text(material.get("content_preview") or material.get("content"))

        for token in re.findall(r"\b[A-Z][a-zA-Z]{3,}(?:\s+[A-Z][a-zA-Z]{3,})?\b", content):

            token = clean_text(token)

            if token and token.lower() != subject_name.lower() and token not in seeds:

                seeds.append(token)

            if len(seeds) >= 28:

                return seeds



    if not seeds:

        seeds = [subject_name, f"Core {subject_name} concepts", f"Applied {subject_name} skills"]

    return seeds





def build_prompt(payload):

    retrieval_context = payload.get("retrieval_context", [])

    context_lines = []

    for index, chunk in enumerate(retrieval_context, start=1):

        context_lines.append(

            f"CONTEXT {index} | {chunk.get('title', 'Material')} | {chunk.get('material_type', 'reference')}\n{clean_text(chunk.get('text', ''))}"

        )



    subject_rules = payload.get("subject_rules", {})

    requested_sections = []

    if payload.get("include_objective"):

        requested_sections.append(

            f"Objective section: {int(payload.get('obj_count', 20))} questions, {int(payload.get('obj_marks', 2))} marks each."

        )

    if payload.get("include_subjective"):

        requested_sections.append(

            f"Theory section: {int(payload.get('subj_count', 5))} questions, {int(payload.get('subj_marks', 8))} marks each."

        )

    if payload.get("include_comprehension"):

        requested_sections.append("Add one comprehension section with one passage and questions derived from it.")



    prompt = f"""

Generate a school examination as strict JSON.



Context:

- Subject: {payload.get('subject_name')}

- Class: {payload.get('class_name')}

- Term: {payload.get('term')}

- Exam title: {payload.get('exam_title')}

- Duration: {payload.get('duration')} minutes

- Total marks target: {payload.get('total_marks')}

- Exam type: {payload.get('exam_type')}



Teacher instructions:

{payload.get('ai_notes') or 'None'}



Topics to emphasize:

{payload.get('topics') or 'None'}



Sample question patterns / past question style to imitate:

{payload.get('sample_questions') or 'None'}



Subject-specific rules:

{os.linesep.join(subject_rules.get('subject', [])) or 'Use balanced subject-appropriate wording.'}



General generation rules:

{os.linesep.join(subject_rules.get('general', [])) or 'Keep the paper balanced.'}



Requested structure:

{os.linesep.join(requested_sections) or 'At least one theory section.'}



Top retrieved learning context:

{os.linesep.join(context_lines) or 'No retrieved context supplied.'}



Return strict JSON with this shape:

{{

  "title": "string",

  "subject": "string",

  "class": "string",

  "time": "string",

  "instructions": "string",

  "sections": [

    {{

      "section": "A",

      "title": "string",

      "instructions": "string",

      "type": "objective|subjective|comprehension",

      "passage": "optional for comprehension",

      "questions": [

        {{

          "number": 1,

          "text": "string",

          "marks": 2,

          "options": {{"A":"", "B":"", "C":"", "D":""}}

        }}

      ]

    }}

  ]

}}



Rules:

- Base the questions primarily on the retrieved context and teacher instructions.

- Spread questions across different ideas instead of repeating one stem style.

- Keep objective distractors plausible and non-duplicated.

- Use wording that is classroom-ready and printable.

- For objective questions include `correct_option` using A-D.`n- For theory/comprehension include `model_answer` and a short `rubric` array.`n- Do not wrap the JSON in markdown.

""".strip()

    return prompt





def call_openai(payload):

    api_key = payload.get("openai_api_key") or os.getenv("OPENAI_API_KEY")

    if not api_key:

        return None, "missing_api_key"



    body = {

        "model": payload.get("model", "gpt-4o-mini"),

        "messages": [

            {"role": "system", "content": "You create balanced school exam papers and return only valid JSON."},

            {"role": "user", "content": build_prompt(payload)},

        ],

        "temperature": 0.4,

        "max_tokens": 3500,

        "response_format": {"type": "json_object"},

    }



    request = urllib.request.Request(

        "https://api.openai.com/v1/chat/completions",

        data=json.dumps(body).encode("utf-8"),

        headers={"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"},

        method="POST",

    )



    try:

        with urllib.request.urlopen(request, timeout=90) as response:

            raw = json.loads(response.read().decode("utf-8"))

            content = raw["choices"][0]["message"]["content"]

            return json.loads(content), None

    except (urllib.error.URLError, urllib.error.HTTPError, KeyError, json.JSONDecodeError) as exc:

        return None, str(exc)





def fallback_exam(payload):

    subject = payload.get("subject_name") or "General Studies"

    class_name = payload.get("class_name") or "Class"

    explicit_topics = split_topics(payload.get("topics", ""))

    sample_questions = payload.get("sample_questions", "")

    materials = payload.get("materials", [])

    retrieval_context = payload.get("retrieval_context", [])

    subject_profile = build_subject_profiles(subject)

    seeds = extract_seed_terms(materials, retrieval_context, sample_questions, explicit_topics, subject)

    random.shuffle(seeds)



    exam = {

        "title": payload.get("exam_title") or f"{subject} Examination",

        "subject": subject,

        "class": class_name,

        "time": f"{payload.get('duration', 60)} minutes",

        "instructions": "Answer all questions. Read each question carefully before answering.",

        "sections": [],

    }



    if payload.get("include_objective"):

        questions = []

        templates = subject_profile["objective_templates"]

        for index in range(1, int(payload.get("obj_count", 20)) + 1):

            topic = seeds[(index - 1) % len(seeds)]

            stem = templates[(index - 1) % len(templates)].format(topic=topic)

            questions.append(

                {

                    "number": index,

                    "text": stem,

                    "marks": int(payload.get("obj_marks", 2)),

                    "options": {

                        "A": f"An accurate idea about {topic}",

                        "B": f"A partly correct detail about {topic}",

                        "C": f"A common misunderstanding about {topic}",

                        "D": f"An unrelated or weak explanation of {topic}",

                    },

                }

            )

        exam["sections"].append({

            "section": "A",

            "title": "Objective Test",

            "instructions": "Choose the best answer for each question.",

            "type": "objective",

            "questions": questions,

        })



    if payload.get("include_subjective"):

        questions = []

        templates = subject_profile["subjective_templates"]

        for index in range(1, int(payload.get("subj_count", 5)) + 1):

            topic = seeds[(index * 2 - 1) % len(seeds)]

            questions.append(

                {

                    "number": index,

                    "text": templates[(index - 1) % len(templates)].format(topic=topic),

                    "marks": int(payload.get("subj_marks", 8)),

                    "model_answer": f"A strong answer should address {topic} accurately with suitable explanation.",

                    "rubric": [

                        "Credit correct explanation of {topic}.".format(topic=topic),

                        "Reward relevant examples and clear presentation.",

                    ],

                }

            )

        exam["sections"].append({

            "section": chr(65 + len(exam["sections"])),

            "title": "Theory / Structured Questions",

            "instructions": "Answer all questions in this section.",

            "type": "subjective",

            "questions": questions,

        })



    if payload.get("include_comprehension"):

        topic = seeds[0]

        exam["sections"].append({

            "section": chr(65 + len(exam["sections"])),

            "title": "Comprehension",

            "instructions": "Read the passage carefully and answer the questions that follow.",

            "type": "comprehension",

            "passage": f"{topic} is an important part of {subject}. This passage focuses on {subject_profile['passage_style']} so learners can connect ideas to classwork and revision.",

            "questions": [

                {"number": 1, "text": "State the main idea of the passage.", "marks": 2, "model_answer": "The answer should identify the central message of the passage.", "rubric": ["Credit the central idea stated clearly."]},

                {"number": 2, "text": f"Explain one reason {topic} matters in {subject}.", "marks": 4},

                {"number": 3, "text": "Give one useful point from the passage for revision.", "marks": 2},

                {"number": 4, "text": "Identify one idea that connects the passage to classroom work.", "marks": 2},

                {"number": 5, "text": "Suggest another suitable title for the passage.", "marks": 2},

            ],

        })



    return exam





def normalize_exam(exam, payload):

    exam.setdefault("title", payload.get("exam_title") or "Generated Examination")

    exam.setdefault("subject", payload.get("subject_name") or "Subject")

    exam.setdefault("class", payload.get("class_name") or "Class")

    exam.setdefault("time", f"{payload.get('duration', 60)} minutes")

    exam.setdefault("instructions", "Answer all questions.")

    exam.setdefault("sections", [])



    for section_index, section in enumerate(exam["sections"]):

        section.setdefault("section", chr(65 + section_index))

        section.setdefault("title", f"Section {chr(65 + section_index)}")

        section.setdefault("instructions", "")

        section.setdefault("type", "subjective")

        section.setdefault("questions", [])

        if section["type"] == "comprehension":

            section["passage"] = clean_text(section.get("passage", ""))

        for question_index, question in enumerate(section["questions"]):

            question["number"] = int(question.get("number", question_index + 1))

            question["text"] = clean_text(question.get("text", ""))

            question["marks"] = max(1, int(question.get("marks", 1)))

            if section["type"] == "objective":

                options = question.get("options") or {}

                question["options"] = {key: clean_text(options.get(key, "")) for key in ["A", "B", "C", "D"]}



    return exam





def main():

    payload = json.loads(sys.stdin.read() or "{}")

    exam, error = call_openai(payload)

    generator = "python_ai"

    warnings = []



    if exam is None:

        exam = fallback_exam(payload)

        generator = "python_fallback"

        if error:

            warnings.append(error)



    exam = normalize_exam(exam, payload)

    sys.stdout.write(json.dumps({"success": True, "exam": exam, "generator": generator, "warnings": warnings}))





if __name__ == "__main__":

    main()

