import re
from typing import List, Dict
from fastapi import HTTPException
import logging
import json
import random
import math
import os 
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from core import Models, StudentNationality
from services.pgvector_service import PGVectorService
from services.openai_service import OpenAIService
from services.chat_database_service import ChatDatabaseService



logger = logging.getLogger(__name__)
class MCQService:
    def __init__(self, pgvector: PGVectorService, db_vector: ChatDatabaseService):
        self.openai_service = OpenAIService()
        if not self.openai_service.is_available():
            logger.warning("Warning: OPENAI_API_KEY not found. MCQ service will be disabled.")
    
        self.pgvector = pgvector
        self.db_service = db_vector

    def _extract_grade_integer(self, grade_str: str) -> int:
        """
        Safely extracts the first integer from a grade string like '4th Grade'.
        This acts as a translator between the string-based MCQ schema and the
        integer-based vector DB schema.
        """
        if not isinstance(grade_str, str):
             raise ValueError(f"Grade must be a string, but got {type(grade_str)}")
        
        numbers = re.findall(r'\d+', grade_str)
        if numbers:
            return int(numbers[0])
        
        # If no numbers are found, we cannot query the vector DB. This is an invalid input.
        raise ValueError(f"Could not extract a numeric grade from the input string: '{grade_str}'")
    



    def generate_and_store_mcqs(
        self, curriculum: str, grade: str, subject: str, unit: str, concept: str, 
        is_arabic: bool, num_questions: int = 5
    ) -> List[Dict]:
        """
        Generates NEW, UNIQUE MCQs with balanced difficulty and Bloom's taxonomy levels.
        Each returned question includes:
        - difficulty_level: 1–10
        - blooms_level: One of ["remember", "understand", "apply", "analysis", "evaluate", "create"]
        """

        if not self.pgvector:
            raise HTTPException(status_code=503, detail="Vector service is not available for context retrieval.")

        logger.info(f"Checking for existing questions for: {curriculum}/{grade}/{subject}/{unit}/{concept}")
        existing_questions = self.pgvector.get_mcqs(
            curriculum, grade, subject, unit, concept, is_arabic, limit=None
        )
        existing_questions_text = "No existing questions found."
        if existing_questions:
            q_list = [f"- {q['question_text']}" for q in existing_questions]
            existing_questions_text = "\n".join(q_list)

        # --- STEP 2: CONTEXT RETRIEVAL ---
        search_query = f"summary of {concept} in {unit} for {subject}"
        query_embedding = self.openai_service.generate_embedding(search_query)
        try:
            grade_for_search = self._extract_grade_integer(grade)
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))
        context_chunks = self.pgvector.search_filtered_nearest(
            query_embedding, grade_for_search, subject, is_arabic, limit=10
        )
        if not context_chunks:
            raise HTTPException(status_code=404, detail="No curriculum context found for this topic.")
        full_context = "\n---\n".join([chunk['chunk_text'] for chunk in context_chunks])

        # --- STEP 3: PROMPT CONSTRUCTION ---
        if is_arabic:
            prompt = f"""
    أنت خبير في تطوير المناهج التعليمية، ومهمتك هي إنشاء **أسئلة اختيار من متعدد جديدة بالكامل**.

    هذه هي الأسئلة الموجودة بالفعل لمفهوم "{concept}":
    --- الأسئلة الموجودة ---
    {existing_questions_text}
    --- نهاية الأسئلة الموجودة ---

    استنادًا فقط إلى المعلومات التالية:
    --- السياق ---
    {full_context}
    --- نهاية السياق ---

    قم بإنشاء {num_questions} سؤالًا جديدًا تمامًا من نوع الاختيار من متعدد (MCQ)، **مختلفة كليًا عن الأسئلة الموجودة أعلاه**.

    ⚠️ **مهم جدًا**:
    يجب أن تشمل الأسئلة مستويات متنوعة من الصعوبة وفق التوزيع التالي:
    - ٤٠٪ أسئلة سهلة (١ إلى ٤)
    - ٣٠٪ أسئلة متوسطة (٥ إلى ٧)
    - ٣٠٪ أسئلة صعبة (٨ إلى ١٠)

    كما يجب أن تغطي الأسئلة مستويات تصنيف بلوم الستة التالية بشكل متوازن تقريبًا:
    - "تذكر" (remember)
    - "فهم" (understand)
    - "تطبيق" (apply)
    - "تحليل" (analysis)
    - "تقييم" (evaluate)
    - "إبداع" (create)

    **صيغة الإخراج** يجب أن تكون مصفوفة JSON صالحة (JSON array) تحتوي على كائنات (objects) بالمفاتيح التالية تمامًا:

    - "question_text": نص السؤال.
    - "difficulty_level": رقم من ١ إلى ١٠.
    - "blooms_level": واحدة من ["remember", "understand", "apply", "analysis", "evaluate", "create"].
    - "question_type": نوع السؤال (مثلاً: "multiple_choice").
    - "correct_answer": الإجابة الصحيحة.
    - "wrong_answer_1" إلى "wrong_answer_4": إجابات خاطئة معقولة.
    - "hint": تلميح للطالب.
    - "question_image_url", "correct_image_url", "wrong_image_url_1" إلى "_4": اتركها كسلسلة فارغة "".

    لا تكتب أي نص خارج مصفوفة JSON.
    """
        else:
            prompt = f"""
    You are an expert curriculum developer. Your task is to generate **entirely new multiple-choice questions (MCQs)** that do NOT overlap with any existing ones.

    Here are the questions that ALREADY EXIST for the concept "{concept}":
    --- EXISTING QUESTIONS ---
    {existing_questions_text}
    --- END EXISTING QUESTIONS ---

    Based ONLY on the following context:
    --- CONTEXT ---
    {full_context}
    --- END CONTEXT ---

    Generate {num_questions} NEW and COMPLETELY DIFFERENT multiple-choice questions.

    ⚠️ **Important Requirements**:
    - Distribute difficulty levels approximately as follows:
    - 40% easy (difficulty 1–4)
    - 30% medium (difficulty 5–7)
    - 30% hard (difficulty 8–10)
    - Also, balance across Bloom's taxonomy levels:
    - "remember"
    - "understand"
    - "apply"
    - "analysis"
    - "evaluate"
    - "create"

    Your response MUST be a valid JSON array of objects.  
    Each object must have **exactly** these keys:

    - "question_text": The text of the question.
    - "difficulty_level": Integer 1–10.
    - "blooms_level": One of ["remember", "understand", "apply", "analysis", "evaluate", "create"].
    - "question_type": The type (e.g., "multiple_choice").
    - "correct_answer": The single correct answer.
    - "wrong_answer_1" to "wrong_answer_4": Plausible incorrect answers.
    - "hint": Helpful explanation or guidance.
    - "question_image_url": ""
    - "correct_image_url": ""
    - "wrong_image_url_1" to "wrong_image_url_4": ""

    Do not include any text outside the JSON array.
    """

        # --- STEP 4: CALL LLM ---
        try:
            response = self.openai_service.client.chat.completions.create(
                model=Models.chat,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                response_format={"type": "json_object"}
            )
            response_content = response.choices[0].message.content
            json_response = json.loads(response_content)

            # +++ THIS IS THE NEW, MORE ROBUST PARSING LOGIC +++
            generated_questions = []
            if isinstance(json_response, list):
                # Case 1: The root of the JSON is already a list of questions.
                generated_questions = json_response
            elif isinstance(json_response, dict):
                # Case 2: The root is a dictionary.
                # First, try to find a list within the dictionary's values.
                found_list = next((v for v in json_response.values() if isinstance(v, list)), None)
                if found_list:
                    generated_questions = found_list
                # If no list is found, maybe the dictionary ITSELF is the single question.
                elif "question_text" in json_response:
                    generated_questions = [json_response] # Wrap the single object in a list.
            
            if not generated_questions:
                # If we still have nothing, the format is truly unknown.
                raise ValueError("LLM response did not contain a recognizable question list or object.")
            
        except (json.JSONDecodeError, ValueError, KeyError, StopIteration) as e:
            logger.error(f"Failed to parse MCQ response from LLM: {e}\nRaw Response: {response_content}")
            raise HTTPException(status_code=500, detail="Failed to generate or parse MCQs from AI.")

        # --- STEP 5: STORE ---
        mcqs_to_store = []
        for q in generated_questions:
            mcqs_to_store.append({
                "curriculum": curriculum,
                "grade": grade,
                "subject": subject,
                "unit": unit,
                "concept": concept,
                "is_arabic": is_arabic,
                "difficulty_level": q.get("difficulty_level"),
                "blooms_level": q.get("blooms_level"),
                "question_text": q.get("question_text"),
                "question_type": q.get("question_type", "multiple_choice"),
                "correct_answer": q.get("correct_answer"),
                "wrong_answer_1": q.get("wrong_answer_1"),
                "wrong_answer_2": q.get("wrong_answer_2"),
                "wrong_answer_3": q.get("wrong_answer_3"),
                "wrong_answer_4": q.get("wrong_answer_4"),
                "hint": q.get("hint"),
                "question_image_url": q.get("question_image_url", ""),
                "correct_image_url": q.get("correct_image_url", ""),
                "wrong_image_url_1": q.get("wrong_image_url_1", ""),
                "wrong_image_url_2": q.get("wrong_image_url_2", ""),
                "wrong_image_url_3": q.get("wrong_image_url_3", ""),
                "wrong_image_url_4": q.get("wrong_image_url_4", ""),
            })

        self.pgvector.insert_mcqs(mcqs_to_store)
        return mcqs_to_store


    

    def handle_ask_for_question(self, student_id: str) -> Dict:
        """
        Handles when a student asks for a question. It generates one new question,
        uses an LLM to find a small subset of RELEVANT questions, and then
        RANDOMLY selects one from that subset. This version correctly handles cases
        with a small number of available questions.
        """
        logger.info(f"Handling 'ask_for_question' request for student {student_id}.")
        
        # 1. Get student info and determine topic (No changes here)
        student_info = self.db_service.get_student_info(student_id)
        if not student_info: raise HTTPException(status_code=404, detail="Student not found.")
        grade, is_arabic, subject = student_info['grade'], student_info['is_arabic'], "Science"
        grade_str = f"{grade}th Grade"
        nationality = student_info['nationality']
        if nationality == StudentNationality.EGYPTIAN:
            curriculum = "EGYPTIAN National"
        else:
            curriculum = "SAUDI National"

        recent_history = self.db_service.get_chat_history(student_id, limit=6)
        if not recent_history: raise HTTPException(status_code=400, detail="Cannot ask a question without conversation context.")

        history_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in recent_history])
        topic_prompt = f"""
        Based on the recent conversation below, identify the specific Unit and Concept the student is currently discussing.
        Your response MUST be a valid JSON object with the keys "unit" and "concept".
        Conversation:\n{history_text}
        """
        try:
            response = self.openai_service.client.chat.completions.create(
                model=Models.classification, messages=[{"role": "user", "content": topic_prompt}],
                temperature=0, response_format={"type": "json_object"}
            )
            topic_data = json.loads(response.choices[0].message.content)
            unit, concept = topic_data['unit'], topic_data['concept']
            logger.info(f"Determined current topic for question: Unit='{unit}', Concept='{concept}'")
        except (json.JSONDecodeError, KeyError) as e:
            logger.error(f"Could not determine topic for student {student_id}: {e}")
            raise HTTPException(status_code=500, detail="Could not determine the current topic.")

        # 2. Generate one new question to enrich the pool (No changes here)
        try:
            self.generate_and_store_mcqs(curriculum, grade_str, subject, unit, concept, is_arabic, num_questions=1)
        except Exception as e:
            logger.warning(f"Non-critical error: Failed to generate a new background MCQ: {e}")

        # 3. Retrieve and filter the pool of available questions (No changes here)
        all_mcqs = self.pgvector.get_mcqs(curriculum, grade_str, subject, unit, concept, is_arabic, limit=None)
        if not all_mcqs: raise HTTPException(status_code=404, detail="No questions found for the current topic.")

        asked_question_texts = {msg['content'] for msg in recent_history if msg['role'] == 'assistant'}
        unasked_mcqs = [mcq for mcq in all_mcqs if mcq['question_text'] not in asked_question_texts]
        
        if not unasked_mcqs:
            logger.warning(f"All questions for '{concept}' have been asked recently. Re-using full list.")
            unasked_mcqs = all_mcqs

        # --- THIS IS THE ROBUST TWO-STEP SELECTION LOGIC ---

        # 4. STEP 1 (Filter with AI): Get a SUBSET of relevant questions.
        relevant_question_texts = []
        last_user_message = recent_history[-1]['content']
        
        # Dynamically determine how many questions to ask for.
        # Ask for up to 3, but no more than the number of available questions.
        num_to_select = min(3, len(unasked_mcqs))
        
        # If there's only one question, we don't need to ask the LLM to choose.
        if num_to_select == 1:
            relevant_question_texts = [unasked_mcqs[0]['question_text']]
            logger.info("Only one un-asked question available, selecting it directly.")
        elif num_to_select > 1:
            selection_prompt = f"""
            A student just said: "{last_user_message}"
            Here is a list of available questions about the topic '{concept}':
            {json.dumps([q['question_text'] for q in unasked_mcqs], ensure_ascii=False, indent=2)}

            From the list above, select the {num_to_select} questions that are MOST RELEVANT to what the student just said.

            Your response MUST be a valid JSON object with a single key "relevant_questions" which is an array of the chosen question strings.
            Example: {{"relevant_questions": ["Question text 1", "Question text 2"]}}
            """
            
            try:
                response = self.openai_service.client.chat.completions.create(
                    model=Models.classification,
                    messages=[{"role": "user", "content": selection_prompt}],
                    temperature=0.1,
                    response_format={"type": "json_object"}
                )
                response_data = json.loads(response.choices[0].message.content)
                relevant_question_texts = response_data.get("relevant_questions", [])
                logger.info(f"LLM identified {len(relevant_question_texts)} relevant questions.")
            except Exception as e:
                logger.warning(f"LLM failed to select a relevant subset of questions: {e}. Will select from all available questions.")
        
        # Robust Fallback: If the LLM fails or returns an empty list, use all un-asked questions as the pool.
        if not relevant_question_texts:
            relevant_question_texts = [mcq['question_text'] for mcq in unasked_mcqs]

        # 5. STEP 2 (Select with Randomness): Randomly choose from the relevant subset.
        chosen_question_text = random.choice(relevant_question_texts)
        
        # 6. Find the full MCQ object for the chosen text and return it.
        chosen_mcq = None
        for mcq in unasked_mcqs:
            if mcq['question_text'] == chosen_question_text:
                chosen_mcq = mcq
                break
        
        # Fallback in case the chosen text somehow doesn't match
        if not chosen_mcq:
            chosen_mcq = random.choice(unasked_mcqs)

        logger.info(f"Selected question for student {student_id}: '{chosen_mcq['question_text']}'")
        
        # Add the chosen question's text to history to prevent immediate re-asking
        self.db_service.add_message(student_id, 'assistant', chosen_mcq['question_text'])
        
        return chosen_mcq
    

    def get_dynamic_quiz(
        self, curriculum: str, grade: str, subject: str, unit: str, concept: str, is_arabic: bool, count: int
    ) -> List[Dict]:
        """
        Generates a dynamic quiz. Handles "All" for unit or concept by recursively 
        calling itself for each sub-topic and dividing the question count.
        """
        if not self.pgvector:
            raise HTTPException(status_code=503, detail="Vector service is not available for this feature.")

        # --- RECURSIVE AGGREGATION LOGIC ---

        # Case 1: Broadest scope - All Units in a Subject
        if unit == "All":
            logger.info(f"Broad scope: All Units for Subject '{subject}'. Fetching units...")
            units = self.pgvector.get_distinct_units_from_structure(curriculum, grade, subject)
            if not units:
                raise HTTPException(status_code=404, detail=f"No units found for {subject}.")
            
            final_quiz = []
            num_parts = len(units)
            base_count = count // num_parts
            remainder = count % num_parts

            for i, u in enumerate(units):
                q_count = base_count + (1 if i < remainder else 0)
                if q_count > 0:
                    # Recursive call for each unit, passing "All" for concept
                    logger.info(f"Fetching {q_count} questions for Unit '{u}'...")
                    final_quiz.extend(self.get_dynamic_quiz(
                        curriculum, grade, subject, u, "All", is_arabic, q_count
                    ))
            
            random.shuffle(final_quiz)
            return final_quiz[:count]

        # Case 2: Medium scope - All Concepts in a Unit
        elif concept == "All":
            logger.info(f"Medium scope: All Concepts for Unit '{unit}'. Fetching concepts...")
            concepts = self.pgvector.get_distinct_concepts_from_structure(curriculum, grade, subject, unit)
            if not concepts:
                raise HTTPException(status_code=404, detail=f"No concepts found for {unit}.")

            final_quiz = []
            num_parts = len(concepts)
            base_count = count // num_parts
            remainder = count % num_parts

            for i, c in enumerate(concepts):
                q_count = base_count + (1 if i < remainder else 0)
                if q_count > 0:
                    # Recursive call for each concept (this will hit the base case below)
                    logger.info(f"Fetching {q_count} questions for Concept '{c}'...")
                    final_quiz.extend(self.get_dynamic_quiz(
                        curriculum, grade, subject, unit, c, is_arabic, q_count
                    ))

            random.shuffle(final_quiz)
            return final_quiz[:count]

        # --- BASE CASE: A SINGLE, SPECIFIC CONCEPT ---
        # This is the original logic you wanted to keep.
        else:
            logger.info(f"Base Case: Fetching {count} questions for specific Concept '{concept}'.")
            
            MAX_QUESTIONS_PER_BATCH = 10
            # Generate a proportional number of freshness questions
            num_fresh_questions = min(max(1, math.floor(count / 3)), 5) if count > 0 else 0
            if num_fresh_questions > 0:
                logger.info(f"Generating {num_fresh_questions} new 'freshness' questions.")
                try:
                    self.generate_and_store_mcqs(
                        curriculum=curriculum, grade=grade, subject=subject, unit=unit, concept=concept,
                        is_arabic=is_arabic, num_questions=num_fresh_questions
                    )
                except Exception as e:
                    logger.warning(f"Could not generate 'freshness' questions for the quiz due to an error: {e}")

            # Fetch all available questions for this specific concept
            final_pool = self.pgvector.get_mcqs(
                curriculum=curriculum, grade=grade, subject=subject, unit=unit, concept=concept,
                is_arabic=is_arabic, limit=None
            )

            if not final_pool:
                # If no questions exist at all, we can't proceed for this part.
                logger.warning(f"No questions could be found or generated for '{concept}'. Returning empty list for this part.")
                return []
            
            # If we still don't have enough after freshness, generate more in batches.
            questions_still_needed = count - len(final_pool)
            if questions_still_needed > 0:
                logger.info(f"Need to generate {questions_still_needed} more to meet count of {count}.")
                remaining = questions_still_needed
                while remaining > 0:
                    batch_size = min(remaining, MAX_QUESTIONS_PER_BATCH)
                    try:
                        self.generate_and_store_mcqs(
                            curriculum=curriculum, grade=grade, subject=subject, unit=unit, concept=concept,
                            is_arabic=is_arabic, num_questions=batch_size
                        )
                        remaining -= batch_size
                    except Exception as e:
                        logger.error(f"Failed to generate batch of {batch_size} questions: {e}")
                        break # Prevent infinite loop on failure
                
                # Re-fetch the pool after batch generation
                final_pool = self.pgvector.get_mcqs(
                    curriculum=curriculum, grade=grade, subject=subject, unit=unit, concept=concept,
                    is_arabic=is_arabic, limit=None
                )

            random.shuffle(final_pool)
            
            # Return the number of questions requested for this part of the recursion
            return final_pool[:min(count, len(final_pool))]