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 using a 'Generate-Filter-Retry' loop with Vector Semantic Deduplication.
        This ensures zero redundancy without sending conversation history to the LLM.
        """
        if not self.pgvector:
            raise HTTPException(status_code=503, detail="Vector service is not available for context retrieval.")

        # 1. Retrieve a pool of textbook context chunks (fetched once, rotated in loop)
        context_chunks = self._retrieve_context_pool(grade, subject, unit, concept, is_arabic)
        
        accepted_mcqs = []
        attempts = 0
        max_attempts = 3  # Safety brake
        
        # 2. Loop until we have enough UNIQUE questions
        while len(accepted_mcqs) < num_questions and attempts < max_attempts:
            attempts += 1
            questions_needed = num_questions - len(accepted_mcqs)
            
            # Ask for slightly more than needed to buffer against duplicates being filtered out
            # (e.g., if we need 1, ask for 2. If we need 5, ask for 7)
            ask_count = questions_needed + 2

            # A. Context Rotation: Pick random chunks to ensure variety in every loop
            # This prevents the LLM from fixating on the same part of the text.
            if len(context_chunks) >= 3:
                selected_chunks = random.sample(context_chunks, min(len(context_chunks), 4))
            else:
                selected_chunks = context_chunks
            
            current_context_text = "\n---\n".join([c['chunk_text'] for c in selected_chunks])

            # B. Generate Raw Questions from LLM
            # (Uses your detailed prompts regarding Bloom's and Difficulty)
            raw_candidates = self._generate_candidate_batch(
                current_context_text, ask_count, grade, is_arabic
            )

            # C. Filter & Embed: Checks DB for semantic duplicates
            valid_batch = self._filter_and_process_questions(
                raw_candidates, accepted_mcqs, curriculum, grade, subject, unit, concept, is_arabic
            )

            accepted_mcqs.extend(valid_batch)

        # 3. Store Final Result
        if accepted_mcqs:
            # Slice to exact number requested (in case we generated extras)
            final_set = accepted_mcqs[:num_questions]
            self.pgvector.insert_mcqs(final_set)
            for q in final_set:
                q.pop('embedding', None)
            return final_set
        
        # If we failed to generate enough
        if not accepted_mcqs:
            logger.warning(f"Could not generate unique questions for {concept} after {attempts} attempts.")
            return []
        
        
            
        return accepted_mcqs


    def _retrieve_context_pool(self, grade, subject, unit, concept, is_arabic):
        """Fetches a larger pool of relevant textbook content to allow for rotation."""
        search_query = f"summary of {concept} in {unit} for {subject}"
        query_emb = self.openai_service.generate_embedding(search_query)
        
        try:
            grade_int = self._extract_grade_integer(grade)
        except ValueError:
            grade_int = 0 
        
        # Fetch 15 chunks (instead of 10) so we can rotate through them
        chunks = self.pgvector.search_filtered_nearest(
            query_emb, grade_int, subject, is_arabic, limit=15
        )
        if not chunks:
            # Fallback text if DB is empty
            return [{"chunk_text": f"Topic summary: {concept} in {unit}"}]
        return chunks

    def _generate_candidate_batch(self, context_text: str, count: int, grade: str, is_arabic: bool) -> List[Dict]:
        """
        Calls LLM to generate questions using the detailed, original prompts 
        to ensure difficulty balance and Bloom's taxonomy coverage.
        """
        
        if is_arabic:
            prompt = f"""
    أنت خبير في تطوير المناهج التعليمية، ومهمتك هي إنشاء **أسئلة اختيار من متعدد جديدة بالكامل**.

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

    قم بإنشاء {count} سؤالًا جديدًا تمامًا من نوع الاختيار من متعدد (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)**.

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

    Generate {count} NEW 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.
    """

        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)

            # Robust Parsing Logic
            if isinstance(json_response, list):
                return json_response
            elif isinstance(json_response, dict):
                # Try to find a list inside values
                found_list = next((v for v in json_response.values() if isinstance(v, list)), None)
                if found_list: return found_list
                # Maybe the dict itself is the question
                if "question_text" in json_response: return [json_response]
            
            return []

        except Exception as e:
            logger.error(f"Failed to generate batch from LLM: {e}")
            return []

    def _filter_and_process_questions(
        self, candidates: List[Dict], current_batch: List[Dict],
        curriculum, grade, subject, unit, concept, is_arabic
    ) -> List[Dict]:
        """
        1. Generates Embedding for new question.
        2. Checks DB if similar vector exists.
        3. Returns only unique questions with embedding attached.
        """
        valid = []
        for q in candidates:
            q_text = q.get("question_text", "").strip()
            if not q_text: continue

            # 1. Local check (prevent duplicates within the current generation loop)
            # Check against questions accepted in previous loops + valid ones in this loop
            if any(x['question_text'] == q_text for x in current_batch + valid):
                continue

            try:
                # 2. Generate Embedding (This creates the vector)
                q_vector = self.openai_service.generate_embedding(q_text)
                
                # 3. Check DB for semantic duplicate using the vector
                is_duplicate = self.pgvector.check_similarity_existence(
                    vector=q_vector, 
                    curriculum=curriculum, 
                    concept=concept, 
                    threshold=0.92 # 92% similarity means it's a duplicate
                )
                
                if is_duplicate:
                    logger.info(f"Vector duplicate filtered: {q_text[:30]}...")
                    continue

                # 4. Add Metadata & Vector to the object for storage
                q.update({
                    "curriculum": curriculum, "grade": grade, "subject": subject,
                    "unit": unit, "concept": concept, "is_arabic": is_arabic,
                    "embedding": q_vector, # Store vector!
                    "question_type": q.get("question_type", "multiple_choice"),
                    "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", ""),
                })
                valid.append(q)

            except Exception as e:
                logger.error(f"Vector check failed for question '{q_text[:20]}': {e}")
                continue
                
        return valid


    

    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


        # 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))]