import os
import json
import random
from openworm_ai.quiz.QuizModel import MultipleChoiceQuiz, Question, Answer
from openworm_ai.quiz.TemplatesCorpus import TEXT_ANSWER_EXAMPLE
from openworm_ai.utils.llms import ask_question_get_response, LLM_GPT4o

indexing = ["A", "B", "C", "D"]
TOKEN_LIMIT = 30_000  # 🔹 Keeps request within OpenAI's limits

# **STRICT Prompt to prevent external knowledge**
STRICT_GENERATE_Q = """
🔹 **TASK:** Generate exactly <QUESTION_NUMBER> multiple-choice questions using **only** the provided text.  
- The questions must be **highly specific** and **cannot** come from general knowledge.  
- If the topic is not in the provided text, **DO NOT** generate a question about it.  
- Questions should challenge **researchers** and **advanced students**.  
- DO NOT include the sources in the questions (...according to [source])
🔹 **FORMAT:**  
QUESTION: <Insert question>  
CORRECT ANSWER: <Correct answer>  
WRONG ANSWER: <Wrong answer 1>  
WRONG ANSWER: <Wrong answer 2>  
WRONG ANSWER: <Wrong answer 3>  

📌 **IMPORTANT:** If the text does not have enough content for <QUESTION_NUMBER> questions, generate as many as possible.  
"""


def load_limited_documents(file_path, max_tokens=TOKEN_LIMIT, num_chunks=5):
    """Loads a JSON document in chunks to generate questions in batches."""
    if not os.path.exists(file_path):
        print(f"⚠ Warning: {file_path} not found. Exiting...")
        return []

    try:
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        extracted_texts = []
        all_text = ""

        # Extract all sections
        for title, doc_contents in data.items():
            all_text += f"\n📌 **{title}**\n"
            for section, details in doc_contents.get("sections", {}).items():
                all_text += f"🔹 **{section}**:\n"
                if "paragraphs" in details:
                    all_text += (
                        " ".join([p["contents"] for p in details["paragraphs"]])
                        + "\n\n"
                    )

        # Split document into chunks
        words = all_text.split()
        chunk_size = len(words) // num_chunks
        for i in range(num_chunks):
            start = i * chunk_size
            end = (i + 1) * chunk_size if i < num_chunks - 1 else len(words)
            chunk_text = " ".join(words[start:end])
            extracted_texts.append(chunk_text)

        return extracted_texts

    except (json.JSONDecodeError, UnicodeDecodeError, PermissionError) as e:
        print(f"⚠ Error reading {file_path}: {e}")
        return []


def save_quiz(num_questions=100, num_answers=4, llm_ver=LLM_GPT4o, temperature=0):
    """Generates and saves a quiz using GPT-4o while ensuring all content is from documents in batches."""

    num_batches = num_questions // 20  # Generate in batches of 20 questions
    document_chunks = load_limited_documents(
        "processed/json/papers/Corsi_et_al_2015.json", num_chunks=num_batches
    )

    if not document_chunks:
        print("⚠ Error: No valid document chunks found.")
        return

    # Initialize quiz
    quiz = MultipleChoiceQuiz(
        title=f"{llm_ver.replace(':', '_')}_{num_questions}questions_celegans_batched",
        source=f"Generated by {llm_ver}, temperature: {temperature}",
    )

    question_count = 0

    # Loop over document chunks and generate questions in batches
    for i, chunk_text in enumerate(document_chunks):
        print(f"📝 Generating batch {i + 1} (20 questions) from document chunk...")

        # Create strict prompt for the batch
        question_prompt = (
            STRICT_GENERATE_Q.replace("<QUESTION_NUMBER>", str(20))
            + TEXT_ANSWER_EXAMPLE
            + "\n\n🔹 **Use ONLY the following document knowledge for questions:**\n\n"
            + chunk_text
        )

        response = ask_question_get_response(question_prompt, llm_ver, temperature)

        # Ensure GPT-4o generated questions
        questions_generated = response.count("QUESTION:")
        if questions_generated < 15:
            print(
                f"⚠ Warning: GPT-4o generated only {questions_generated} questions in batch {i + 1}. Skipping batch."
            )
            continue

        # Parse and add questions to quiz
        last_question = None
        indexing = ["1", "2", "3", "4"]

        for line in response.split("\n"):
            if question_count >= num_questions:
                break  # Stop at exactly 100 questions

            if len(line.strip()) > 0:
                if "QUESTION" in line or line.strip().endswith("?"):
                    question_text = line.split(":", 1)[-1].strip()
                    print(f"Question: <{question_text}>")
                    last_question = Question(question=question_text)
                    quiz.questions.append(last_question)
                    question_count += 1
                elif "CORRECT ANSWER" in line:
                    correct_ans = line.split(":", 1)[-1].strip()
                    i = len(last_question.answers)
                    last_question.answers.append(Answer(indexing[i], correct_ans, True))
                elif "WRONG ANSWER" in line:
                    wrong_ans = line.split(":", 1)[-1].strip()
                    i = len(last_question.answers)
                    last_question.answers.append(Answer(indexing[i], wrong_ans, False))

        if question_count >= num_questions:
            break  # Stop if we reached 100 questions

    # Ensure quiz has enough valid questions before saving
    if len(quiz.questions) < num_questions * 0.8:
        print(
            "⚠ Error: Not enough valid questions were generated. Quiz will not be saved."
        )
        return

    print("===============================\n  Generated quiz:\n")
    print(quiz.to_yaml())

    quiz.to_json_file(
        f"openworm_ai/quiz/samples/{llm_ver.replace(':', '_')}_{num_questions}questions_celegans_batched.json"
    )


if __name__ == "__main__":
    import sys

    llm_ver = LLM_GPT4o  # Always use GPT-4o
    print(f"Selected LLM: {llm_ver}")

    if "-ask" in sys.argv:
        quiz_json = f"openworm_ai/quiz/samples/{llm_ver.replace(':', '_')}_100_questions_celegans_corpus.json"
        quiz = MultipleChoiceQuiz.from_file(quiz_json)

        total_qs = 0
        total_correct = 0
        wrong_answers = "Incorrect answers:\n"

        for qi, question in enumerate(quiz.questions):
            q = question["question"]

            from openworm_ai.quiz.Templates import ASK_Q

            answers = ""
            random.shuffle(question["answers"])

            presented_answers = {}
            for index, answer in enumerate(question["answers"]):
                ref = indexing[index]
                present = f"{ref}: {answer['ans']}"
                if answer["correct"]:
                    correct_answer = ref
                    correct_text = present
                presented_answers[ref] = present
                answers += f"{present}\n"

            full_question = ASK_Q.replace("<QUESTION>", q).replace("<ANSWERS>", answers)

            resp = ask_question_get_response(
                full_question, llm_ver, print_question=False
            ).strip()

            total_qs += 1
            correct_guess = resp == correct_answer

            print(
                f" >> {qi}) {q} → Guess: {resp}, Correct: {correct_answer} → {correct_guess}"
            )

    else:
        save_quiz(100, 4, llm_ver, temperature=0.2)
