import logging
import os
from typing import List, Dict, Optional
from fastapi import HTTPException
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from core import StudentNationality, Models
from services.pgvector_service import PGVectorService
from services.openai_service import OpenAIService
from services.chat_database_service import ChatDatabaseService, StudyLanguage
from services.pedagogy_service import PedagogyService
from services.connection_pool import ConnectionPool

logger = logging.getLogger(__name__)

# Enhanced system prompts (keeping existing ones but adding curriculum awareness instructions)
ENHANCED_SYSTEM_PROMPTS: Dict[tuple, str] = {
    # ---------- Egyptian + Arabic ----------
    (StudentNationality.EGYPTIAN, StudyLanguage.ARABIC): """
إنك مُدرِّس لطفل في ابتدائي اسمه {student_name} في الصف {grade}.  
فقط لو الطفّل سأل عن هويتك بصراحة ووضح (مثل "إنت مين؟"، "عرّفني بنفسك"، "إنت بتعمل إيه هنا؟")،  
رُد بالنصّ الثابت ده:  
"أنا عَنان مؤسِّس شارع العلوم، وإنت هنا على مَنَصّة Science Street Lab،  
وأنا هنا عشان أساعدك تتعلَّم أي حاجة عايز تتعلَّمها في العلوم."  

لو سأل أسئلة عامة زي "نت عارف انا مين؟" أو "إزيك؟" أو "شكرا"، رد بطريقة طبيعية ودودة باستخدام اسمه {student_name}.

**عن المنهج وطريقة الشرح:**
- إنت عارف المنهج المصري للصف {grade} كويس من ملف الـ JSON.  
- لو الطّفل سأل بشكل عام ("إحنا هناخد إيه؟"، "وريني المنهج"، "إيه المواضيع؟")، اعرضله المنهج بشكل واضح ومنظّم.  
- لو سأل عن حاجة مُعيّنة، استخدم السياق من المنهج واربطها بالوحدات والأفكار اللي ليها علاقة.  
- دايمًا قول مكان الدرس فين: "ده في الوحدة الأولى، المفهوم التاني".  
- وضّح الروابط: "ده ليه علاقة باللي اتعلمناه عن..." أو "ده هيوصلنا للي هنتعلمه بعد كده عن...".  

ملاحظة مُلزمة: كلمة "منصّة" لازم تكتبها دايماً كده بالضبط: **مَنَصّة** (بالفتحة على الميم والنون)،  
عشان الـTTS ينطقها صح.

في باقي الردود، رَد باللهجة المصريّة الطبيعيّة كأنّك بتكَلّم {student_name} قصادك.  
خَلّي الكلام بسيط، واضح، وقَريب من وجدنه.  
الجُملَ قُصيرَة ومُترابطة، مُش مَقطَّعة.  
اشرح كأنّك بتحكي له حكاية أو بتوريّه حاجَة من الحَياة حوالينا، مُش بتقرا من كتاب.  
مُمكن تُذكر اسم {student_name} مَرّة واحدة في أوّل الرّد فَقَط.  
بعد كدا مَمنوع تكرار الاسم في نَفس الرّد، حَتّى في الأسئلة الختاميّة.  
مَمنوع تستَعملُ أي ألقاب زي "يا بَطَل" أو "يا شاطر"، الاسم الأوَّل بَس.  
ولو الرّد قُصيرَ جدّاً (جملة أو اتنين)، مُمكن تستَغنَى عن ذكر الاسم خالص.  
لو فيه مُصطَلَح صَعب، فَسّره بكلمة أسهَل.  
لو فيه رَمز كيمياوي زي H2O أو CO2، اكتبه زي ما هو.  
الأرقام العاديّة اكتبها بالحروف العربي زي "اتنين" أو "تَلاتة".  
استخدمُ التشكيل الكامل على كُلّ الكلام عَشان يطّلع بالصّوت زي نُطق اللّهجة المصريّة الطَبيعيّ.  
لو {student_name} مكتوب بالإنجليزي، اكتبه دايماً بالعَربي في ردودك.  

لَمّا تُذكر الصف {grade}، قُله بالطريقة الطبيعيّة زي ما الأطفال بيقولوها: الصف 4 = سنة رابعة ابتدائي، الصف 5 = سنة خامسة ابتدائي، وهكذا.  

الهَدف: رَد قُصيرَ يُعلِّم ويُوصَّل المَعلومة، ويُبان إن "عَنان" بيشرَح للطفل جوّه مَنَصّة "شارِع العلوم"، مُش كتاب بيتقري.
""",

    # ---------- Saudi + Arabic ----------
    (StudentNationality.SAUDI, StudyLanguage.ARABIC): """
إنت مُدرِّس لطفل في ابتدائي اسمه {student_name} في الصف {grade}.  
فقط لو الطفل سأل عن هويتك بصراحة ووضح (مثل "إنت مين؟"، "عرِّفني بنفسك"، "إنت وش تسوي هنا؟")،  
رُد بالنص الثابت هذا:  
"أنا عَنان مؤسِّس شارع العلوم، وإنت هنا على مَنَصّة Science Street Lab،  
وأنا هنا عشان أساعدك تتعلَّم أي حاجة عايز تتعلَّمها في العلوم."  

لو سأل أسئلة عامة مثل "نت عارف انا مين؟" أو "كيفك؟" أو "شكرا"، رد بطريقة طبيعية ودودة باستخدام اسمه {student_name}.

**عن المنهج وطريقة الشرح:**
- إنت تعرف المنهج السعودي للصف {grade} مضبوط من ملف الـ JSON.  
- لو الولد سأل بشكل عام ("وش بندرس؟"، "ورني المنهج"، "إيش المواضيع؟")، عرّض له المنهج بشكل واضح ومرتب.  
- لو سأل عن موضوع محدد، استخدم السياق من المنهج واربطه بالوحدات والأفكار اللي لها علاقة.  
- دايمًا قول وين مكان الدرس: "هذا في الوحدة الأولى، المفهوم الثاني".  
- وضّح الترابط: "هذا مرتبط باللي تعلمناه عن..." أو "هذا يودينا للي راح نتعلمه بعدين عن...".  

ملاحظة مُلزمة: كلمة "منصّة" لازم تكتبها دايماً كده بالضبط: **مَنَصّة** (بالفتحة على الميم والنون)،  
عشان الـTTS ينطقها صح.

في باقي الردود، رَد باللهجة السعوديّة الطبيعيّة، كأنك تشرح له قدّامك.  
خل الشرح واضح وسهل، لكن لا يكون ناشف.  
اشرح كأنك تسولف معه وتشبّه بأشياء من حياته اليوميّة.  

اذكر اسم {student_name} مرّة وحدة فقط في بداية الرد.  
بعد كذا لا تكرره في النص ولا في الأسئلة الختاميّة.  
ممنوع تستخدم أي ألقاب مثل "يا بطل" أو "يا شاطر"، الاسم الأول يكفي.  
ولو الرد قصير جداً (جملة أو جملتين), تقدر ما تذكر الاسم أبداً.  

لو فيه مصطلح صعب، فسِّره بكلمة أبسط.  
الرموز الكيمياوية مثل H2O أو CO2 تكتب مثل ما هي.  
الأرقام في الكلام العادي تكتبها بالحروف العربي زي "اثنين" أو "ثلاثة".  
استخدم التشكيل بس على الكلمات اللي ممكن الـTTS يخبّص فيها أو يقرأها خطأ، واترك الباقي بدون تشكيل عشان يطلع طبيعي.  
لو {student_name} مكتوب بالإنجليزي، اكتبه دايماً بالعربي في ردودك.  

لما تذكر الصف {grade}، قولها بالطريقة اللي الطفل متعود يسمعها: الصف 4 = رابع ابتدائي، الصف 5 = خامس ابتدائي، وهكذا.  

الهدف: رد مبسِّط، قريب، ويبيِّن إن "عَنان" يشرح للطفل جوّه مَنَصّة "شارع العلوم"، مو يقرأ من كتاب.
""",

    # -------- Egyptian English --------
    (StudentNationality.EGYPTIAN, StudyLanguage.ENGLISH): """
إنت مُدرِّس لطفل في ابتدائي اسمه {student_name} في الصف {grade}.  لو الطفّل سأل عن هويتك بصراحة (زي "إنت مين؟"، "عرِّفني بنفسك")،  
رُد بالنصّ الثابت ده:  
"أنا عَنان مؤسس شارع العلوم، وإنت هنا على مَنَصّة Science Street Lab،  
وأنا هنا عشان أساعدك تتعلَّم أي حاجة عايز تتعلَّمها في العلوم."

لو سأل أسئلة عامة (زي "إزيك؟"، "شكراً")، رد بطريقة طبيعية ودودة باستخدام اسمه {student_name}.  

**عن المنهج وطريقة الشرح:**
- إنت عارف المنهج المصري للصف {grade} كويس من ملف الـ JSON.  
- لو الطّفل سأل بشكل عام ("إحنا هناخد إيه؟"، "وريني المنهج"، "إيه المواضيع؟")، اعرضله المنهج بشكل واضح ومنظّم.  
- لو سأل عن حاجة مُعيّنة، استخدم السياق من المنهج واربطها بالوحدات والأفكار اللي ليها علاقة.  
- دايمًا قول مكان الدرس فين: "ده في الوحدة الأولى، المفهوم التاني".  
- وضّح الروابط: "ده ليه علاقة باللي اتعلمناه عن..." أو "ده هيوصلنا للي هنتعلمه بعد كده عن...".  


بالنسبة لأسئلة العلوم أو المنهج:  
- Always answer **in English first**.  
- After answering, ask: *"اشرحهالك بالعربي اوّ بشكل ابسط؟"*  
- If the child says yes (or asks in Arabic), then give a **mixed explanation**  
  (**English for terminologies + simple Arabic for explanation**).  

احرص إن الشرح يكون بسيط، قصير، واضح، وكأنك بتحكي له من الحياة اليومية.  
اذكر اسم {student_name} مرة واحدة بس في بداية الرد. متكررهوش تاني.  
ممنوع تستخدم ألقاب زي "يا بطل" أو "يا شاطر".  
لو الرد قصير جداً (جملة أو اتنين) ممكن تستغنى عن الاسم.  

لما تذكر الصف {grade}، قولها بالطريقة اللي الأطفال المصريين بيقولوها:  
الصف 4 = سنة رابعة ابتدائي، الصف 5 = سنة خامسة ابتدائي، وهكذا.  

المصطلحات العلميّة: سيبها بالإنجليزي (**roots**, **photosynthesis**, **glucose**) مع شرح بسيط.  
الصيغ الكيمياويّة زي H2O أو CO2 لازم تكتب زي ما هي.  
الأرقام في الجُملَ العاديّة بالإنجليزي بالحروف (two, three).  

الهَدف: إجابة بالإنجليزي واضحة ومبسّطة، وبعدها عرض مساعدة إضافية بالعربي لو الطفّل حب،  
بحيث يبان إن "عَنان" بيشرح جوّه مَنَصّة "شارِع العُلوم".
""",

# -------- Saudi English --------
(StudentNationality.SAUDI, StudyLanguage.ENGLISH): """
إنت مُدرِّس لطفل في ابتدائي اسمه {student_name} في الصف {grade}.  
لو الطفل سأل عن هويتك بصراحة (زي "إنت مين؟"، "عرِّفني بنفسك"، "إنت وش تسوي هنا؟")،  
رُد بالنصّ الثابت هذا:  
"أنا عَنان مؤسس شارع العلوم، وإنت هنا على مَنَصّة Science Street Lab،  
وأنا هنا عشان أساعدك تتعلَّم أي حاجة عايز تتعلَّمها في العلوم."

لو سأل أسئلة عامة (زي "كيفك؟"، "شكراً")، رد بطريقة طبيعية ودودة باستخدام اسمه {student_name}.  

**عن المنهج وطريقة الشرح:**
- إنت تعرف المنهج السعودي للصف {grade} مضبوط من ملف الـ JSON.  
- لو الولد سأل بشكل عام ("وش بندرس؟"، "ورني المنهج"، "إيش المواضيع؟")، عرّض له المنهج بشكل واضح ومرتب.  
- لو سأل عن موضوع محدد، استخدم السياق من المنهج واربطه بالوحدات والأفكار اللي لها علاقة.  
- دايمًا قول وين مكان الدرس: "هذا في الوحدة الأولى، المفهوم الثاني".  
- وضّح الترابط: "هذا مرتبط باللي تعلمناه عن..." أو "هذا يودينا للي راح نتعلمه بعدين عن...".  


بالنسبة لأسئلة العلوم أو المنهج:  
- Always answer **in English first**.  
- After answering, ask: *"اشرحهالك بالعربي اوّ بشكل ابسط؟"* 
- If the child says yes (or asks in Arabic), then give a **mixed explanation**  
  (**English for terminologies + simple Arabic for explanation**).  

خل الشرح واضح وسهل وبأمثلة من حياة الطفل اليوميّة.  
اذكر اسم {student_name} مرّة وحدة فقط في بداية الرد. لا تكرره في نفس الرد.  
ممنوع تستخدم ألقاب زي "يا بطل" أو "يا شاطر". الاسم الأول يكفي.  
ولو الرد قصير جداً (جملة أو جملتين)، ممكن ما تذكر الاسم أبداً.  

لما تذكر الصف {grade}، قولها بالطريقة اللي الأطفال السعوديين متعودين عليها:  
الصف 4 = رابع ابتدائي، الصف 5 = خامس ابتدائي، وهكذا.  

المصطلحات العلميّة: خليها بالإنجليزي (**roots**, **photosynthesis**, **glucose**) مع شرح مبسّط.  
الصيغ الكيمياويّة مثل H2O أو CO2 لازم تكتب مثل ما هي.  
الأرقام في النصوص العاديّة بالإنجليزي بالحروف (two, three).  

الهدف: إجابة بالإنجليزي مبسّطة، وبعدها عرض مساعدة بالعربي لو الطفل حب،  
عشان يبان إن "عَنان" يشرح داخل مَنَصّة "شارع العلوم".
"""

}


class AgentService:
    """Enhanced service class for handling AI agent conversations with JSON-based curriculum structure"""

    def __init__(self, use_pgvector: bool = True, pool_handler: Optional[ConnectionPool] = None):
        self.openai_service = OpenAIService()
        if not self.openai_service.is_available():
            logger.warning("Warning: OPENAI_API_KEY not found. Agent service will be disabled.")
            self.client = None
        else:
            self.client = self.openai_service.client

        self.pool_handler = pool_handler
        if self.pool_handler is None:
            self.pool_handler = ConnectionPool(
                minconn=1,
                maxconn=20,
                dbname=os.getenv("DB_NAME"),
                user=os.getenv("DB_USER"),
                password=os.getenv("DB_PASSWORD"),
                host=os.getenv("DB_HOST"),
                port=os.getenv("DB_PORT")
            )
        
        self.use_pgvector = use_pgvector
        # Pass the same pool handler to both services
        self.db_service = ChatDatabaseService(self.pool_handler)
        if self.use_pgvector:
            self.pgvector = PGVectorService(self.pool_handler)
            # Setup curriculum table if needed
            self.pgvector.setup_curriculum_table()
        else:
            self.pgvector = None

        self.pedagogy_service = PedagogyService()
        self.student_info = {}

    def is_available(self) -> bool:
        return self.client is not None

    def get_conversation_history(self, student_id: str) -> List[Dict[str, str]]:
        """Get conversation history from database"""
        try:
            return self.db_service.get_chat_history(student_id)
        except Exception as e:
            logger.error(f"Error getting conversation history for {student_id}: {e}")
            return []

    def add_message_to_history(self, student_id: str, message: str, role: str = "user"):
        """Add message to database"""
        try:
            self.db_service.add_message(student_id, role, message)
            # Limit history to prevent growth
            self.db_service.limit_history(student_id, max_messages=38)
        except Exception as e:
            logger.error(f"Error adding message to history for {student_id}: {e}")

    def classify_query_type(self, query: str, student_info: Dict) -> str:
        """Enhanced query classification using LLM based on JSON structure"""
        if not self.is_available():
            return "specific_content"
            
        is_arabic = student_info.get('is_arabic', True)
        grade = student_info.get('grade', 4)
        
        classification_prompt = f"""
        صنف السؤال التالي إلى إحدى الفئات الأربع:

        1. "general_chat" - أسئلة دردشة عامة وشخصية عن الطالب أو المدرس
           أمثلة: "إنت مين؟", "إزيك؟", "نت عارف انا مين؟", "انت عارف انا في سنة كام؟", "شكرا ليك", "who are you?", "how are you?", "do you know me?", "thank you"

        2. "overview" - أسئلة عن نظرة عامة على المنهج أو المحتوى الكامل
           أمثلة: "ماذا ندرس؟", "أظهر المنهج", "what do we study?", "show curriculum"
           
        3. "navigation" - أسئلة عن وحدة أو مفهوم معين
           أمثلة: "ما في الوحدة الأولى؟", "what's in unit 1?", "أخبرني عن مفهوم الطاقة"
           
        4. "specific_content" - أسئلة محددة عن موضوع علمي معين
           أمثلة: "ما هو التمثيل الضوئي؟", "what is photosynthesis?", "كيف تعمل الخلية؟"

        السؤال: "{query}"
        الطالب يدرس باللغة: {"العربية" if is_arabic else "الإنجليزية"}
        الصف: {grade}
        
        رد فقط بكلمة واحدة: general_chat أو overview أو navigation أو specific_content
        """
        
        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": classification_prompt}],
                temperature=0,
                max_tokens=10
            )
            classification = response.choices[0].message.content.strip().lower()
            
            if classification in ["general_chat", "overview", "navigation", "specific_content"]:
                logger.info(f"Query classified as: {classification} for query: '{query}'")
                return classification
            else:
                logger.warning(f"Unknown classification: {classification}, defaulting to specific_content")
                return "specific_content"
                
        except Exception as e:
            logger.warning(f"Error in query classification: {e}, defaulting to specific_content")
            return "specific_content"

    def handle_overview_query(self, student_info: Dict, subject: str = "Science") -> str:
        """Handle curriculum overview queries using JSON-based data"""
        if not self.pgvector:
            if student_info['study_language'] == StudyLanguage.ARABIC:
                return f"عذراً، لا يمكنني عرض المنهج حالياً للصف {student_info['grade']}"
            else:
                return f"Sorry, I cannot show the curriculum for Grade {student_info['grade']} right now"
        
        return self.pgvector.get_overview_response(
            student_info['grade'], 
            student_info['is_arabic'], 
            subject
        )

    def handle_general_chat_query(self, query: str, student_info: Dict) -> str:
        """Handle general chat queries using only student information"""
        student_name = student_info.get('student_name', 'الطالب')
        grade = student_info.get('grade', 4)
        nationality = student_info.get('nationality', 'مصري')
        is_arabic = student_info.get('is_arabic', True)
        study_language = student_info['study_language']
        
        # Create a simple context with student info only
        if is_arabic:
            context = f"""
معلومات الطالب:
- الاسم: {student_name}
- الصف: {grade}
- الجنسية: {nationality}
- لغة الدراسة: {"العربية" if is_arabic else "الإنجليزية"}

السؤال: "{query}"

أجب بناء على معلومات الطالب فقط. لا تستخدم أي معلومات من المنهج أو محتوى تعليمي.
إذا سأل الطالب عن هويتك، استخدم الرد المحدد في التعليمات.
إذا سأل عن معلوماته الشخصية، استخدم البيانات المتاحة أعلاه.
كن ودوداً وبسيطاً في الرد.
            """
        else:
            context = f"""
Student Information:
- Name: {student_name}
- Grade: {grade}
- Nationality: {nationality}
- Study Language: {"Arabic" if is_arabic else "English"}

Question: "{query}"

Answer based only on the student's information above. Do not use any curriculum or educational content.
If the student asks about your identity, use the specified response in the instructions.
If they ask about their personal information, use the data available above.
Be friendly and simple in your response.
            """
        
        return context
        """Handle unit/concept navigation queries using JSON structure"""
        if not self.pgvector:
            return self.handle_overview_query(student_info, subject)
            
        return self.pgvector.get_unit_navigation_response(
            query, 
            student_info['grade'], 
            student_info['is_arabic'], 
            subject
        )

    def generate_enhanced_context(self, search_results: List[Dict], student_info: Dict, query_type: str) -> str:
        """Generate enhanced context with JSON-based curriculum structure awareness"""
        if not search_results:
            return ""
        
        is_arabic = student_info['is_arabic']
        study_language = student_info['study_language']
        grade = student_info['grade']
        
        if study_language == StudyLanguage.ENGLISH:
            context_message = f"📚 من المنهج الإنجليزي لمادة العلوم للصف {grade}:\n\n"
        else:
            context_message = f"📚 من المنهج العربي لمادة العلوم للصف {grade}:\n\n"
        
        for result in search_results:
            # Basic information
            unit_info = f"الوحدة: {result['unit']}" if result.get('unit') else ""
            concept_info = f"المفهوم: {result['concept']}" if result.get('concept') else ""
            lesson_info = f"الدرس: {result['lesson']}" if result.get('lesson') else ""
            
            # Build header
            context_parts = [info for info in [unit_info, concept_info, lesson_info] if info]
            if context_parts:
                context_message += f"**{' → '.join(context_parts)}**\n"
            
            # Add content
            context_message += f"{result['chunk_text']}\n"
            
            # Add curriculum context if available
            if 'curriculum_context' in result:
                ctx = result['curriculum_context']
                if ctx.get('navigation_hint'):
                    context_message += f"\n💡 {ctx['navigation_hint']}\n"
                
                if ctx.get('related_concepts') and query_type == "specific_content":
                    related = ', '.join(ctx['related_concepts'][:3])
                    if is_arabic:
                        context_message += f"🔗 مفاهيم ذات صلة: {related}\n"
                    else:
                        context_message += f"🔗 Related concepts: {related}\n"
            
            context_message += "\n---\n\n"
        
        # Add instruction for using the context
        if study_language == StudyLanguage.ENGLISH:
            context_message += f"استخدم هذه المعلومات لتقديم شرح دقيق للطفل. المنهج إنجليزي فاذكر المصطلحات الإنجليزية مع الشرح بالعربي."
        else:
            context_message += f"استخدم هذه المعلومات لتقديم شرح دقيق ومناسب للطفل باستخدام المصطلحات العربية."
            
        return context_message

    def generate_response(
        self,
        user_message: str,
        student_id: str,
        subject: str = "Science",
        model: str = Models.chat,
        temperature: float = 0.3,
        top_k: int = 3
    ) -> str:
        """Enhanced AI response generation with JSON-based curriculum structure awareness"""
        if not self.is_available():
            raise HTTPException(status_code=500, detail="Agent service not available")

        try:
            # Get student info with explicit language
            student_info = self.db_service.get_student_info(student_id)
            if not student_info:
                raise HTTPException(status_code=404, detail=f"Student with ID {student_id} not found")

            # Extract information
            full_name = student_info.get('student_name', 'الطالب')
            student_name = full_name.split()[0] if full_name else "الطالب"
            study_language = student_info['study_language']
            
            # Map nationality
            nationality_lower = student_info['nationality'].lower().strip()
            nationality_mapping = {
                'egyptian': StudentNationality.EGYPTIAN,
                'saudi': StudentNationality.SAUDI
            }
            nationality = nationality_mapping.get(nationality_lower, StudentNationality.EGYPTIAN)

            # Add user message to DB
            self.add_message_to_history(student_id, user_message, "user")
            conversation_history = self.get_conversation_history(student_id)

            # Classify query type based on JSON structure
            query_type = self.classify_query_type(user_message, student_info)
            logger.info(f"Query type: {query_type} for student {student_name} ({study_language.value})")

            # Get appropriate system prompt
            prompt_key = (nationality, study_language)
            base_system_prompt = ENHANCED_SYSTEM_PROMPTS.get(prompt_key, 
                ENHANCED_SYSTEM_PROMPTS.get((StudentNationality.EGYPTIAN, StudyLanguage.ARABIC), ""))
            
            formatted_base_prompt = base_system_prompt.format(
                student_name=student_name,
                grade=student_info['grade']
            )

            # Add Socratic instructions if any
            socratic_instructions = self.pedagogy_service.get_socratic_instructions(
                student_info['grade'], student_info['nationality']
            )
            if socratic_instructions:
                formatted_base_prompt += f"\n\n{socratic_instructions}"

            # Prepare messages
            messages = []
            has_system_message = conversation_history and conversation_history[0].get("role") == "system"
            if not has_system_message:
                messages.append({"role": "system", "content": formatted_base_prompt})
                self.add_message_to_history(student_id, formatted_base_prompt, "system")
            messages.extend(conversation_history)

            # Handle different query types with JSON-based curriculum awareness
            if query_type == "general_chat":
                # Handle general chat with student info only
                chat_context = self.handle_general_chat_query(user_message, student_info)
                messages.append({"role": "system", "content": f"سياق المحادثة العامة:\n{chat_context}"})
                
            elif query_type == "overview":
                # Direct curriculum overview from JSON
                overview_response = self.handle_overview_query(student_info, subject)
                messages.append({"role": "system", "content": f"المنهج الكامل من ملف JSON:\n{overview_response}"})
                
            elif query_type == "navigation":
                # Unit/concept navigation from JSON structure
                navigation_response = self.handle_navigation_query(user_message, student_info, subject)
                messages.append({"role": "system", "content": f"تفاصيل الوحدة/المفهوم من JSON:\n{navigation_response}"})
                
            elif query_type == "specific_content" and self.pgvector:
                # Enhanced content search with JSON-based curriculum context
                try:
                    query_embedding = self.openai_service.generate_embedding(user_message)
                    search_results = self.pgvector.search_with_curriculum_context(
                        query_embedding=query_embedding,
                        grade=student_info['grade'],
                        subject=subject,
                        is_arabic=student_info['is_arabic'],
                        limit=top_k
                    )

                    relevant_results = [r for r in search_results if r['distance'] < 1.3] if search_results else []

                    if relevant_results:
                        enhanced_context = self.generate_enhanced_context(
                            relevant_results, student_info, query_type
                        )
                        messages.append({"role": "system", "content": enhanced_context})
                        logger.info(f"Added enhanced JSON-based context with {len(relevant_results)} chunks")

                except Exception as e:
                    logger.warning(f"Error in enhanced content search: {e}")

            # Generate response
            response = self.client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature
            )
            ai_response = response.choices[0].message.content.strip()
            if not ai_response:
                raise ValueError("Empty response from AI model")

            # Save AI response
            self.add_message_to_history(student_id, ai_response, "assistant")
            
            logger.info(f"Generated {query_type} response for {student_name} ({study_language.value}): {len(ai_response)} characters")
            
            return ai_response

        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"Error generating AI response: {e}")
            raise HTTPException(status_code=500, detail=f"AI response generation failed: {str(e)}")

    def search_similar(self, query_embedding: List[float], student_id: str, 
                      subject: str = "chemistry", top_k: int = 3):
        """Search similar content with student-specific filtering and JSON-based curriculum awareness"""
        if not self.pgvector:
            raise HTTPException(status_code=400, detail="PGVector service not enabled")
        
        try:
            student_info = self.db_service.get_student_info(student_id)
            if not student_info:
                raise HTTPException(status_code=404, detail=f"Student with ID {student_id} not found")
            
            logger.info(f"Enhanced search for student {student_info['student_name']} who studies in {student_info['study_language'].value}")
            
            return self.pgvector.search_with_curriculum_context(
                query_embedding=query_embedding,
                grade=student_info['grade'],
                subject=subject,
                is_arabic=student_info['is_arabic'],
                limit=top_k
            )
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"Error in enhanced search_similar: {e}")
            raise HTTPException(status_code=500, detail=f"Enhanced search failed: {str(e)}")

    def get_available_subjects(self, student_id: str) -> List[str]:
        """Get available subjects for the student based on their grade and language from JSON data"""
        if not self.pgvector:
            return []
        
        try:
            student_info = self.db_service.get_student_info(student_id)
            if not student_info:
                return []
            
            return self.pgvector.get_subjects_by_grade_and_language(
                student_info['grade'], 
                student_info['is_arabic']
            )
        except Exception as e:
            logger.error(f"Error getting available subjects for {student_id}: {e}")
            return []

    def get_curriculum_overview(self, student_id: str, subject: str = "Science") -> Dict:
        """Get curriculum overview for a specific student based on JSON data"""
        if not self.pgvector:
            return {"error": "PGVector service not available"}
            
        try:
            student_info = self.db_service.get_student_info(student_id)
            if not student_info:
                return {"error": "Student not found"}
            
            curriculum = self.pgvector.get_curriculum_structure(
                student_info['grade'],
                student_info['is_arabic'], 
                subject
            )
            
            if not curriculum:
                return {"error": f"No curriculum found for Grade {student_info['grade']} in {student_info['study_language'].value}"}
            
            return {
                "student_info": {
                    "name": student_info['student_name'],
                    "grade": student_info['grade'],
                    "study_language": student_info['study_language'].value,
                    "nationality": student_info['nationality']
                },
                "curriculum": curriculum,
                "available_units": self.pgvector.get_units_for_grade(
                    student_info['grade'], student_info['is_arabic'], subject
                ),
                "source": "JSON-based curriculum structure"
            }
            
        except Exception as e:
            logger.error(f"Error getting curriculum overview: {e}")
            return {"error": str(e)}

    def export_conversation(self, student_id: str) -> List[Dict[str, str]]:
        """Export conversation history for a student"""
        try:
            return self.get_conversation_history(student_id)
        except Exception as e:
            logger.error(f"Error exporting conversation for {student_id}: {e}")
            return []

    def clear_conversation(self, student_id: str) -> Dict[str, str]:
        """Clear conversation history for a student"""
        try:
            self.db_service.clear_history(student_id)
            return {
                "status": "success",
                "message": f"Conversation cleared for student {student_id}"
            }
        except Exception as e:
            logger.error(f"Error clearing conversation: {e}")
            return {
                "status": "error",
                "message": f"Failed to clear conversation: {str(e)}"
            }

    def get_agent_stats(self, student_id: str) -> Dict:
        """Get conversation statistics for a student with enhanced JSON-based curriculum info"""
        try:
            history = self.get_conversation_history(student_id)
            user_messages = [msg for msg in history if msg['role'] == 'user']
            assistant_messages = [msg for msg in history if msg['role'] == 'assistant']
            system_messages = [msg for msg in history if msg['role'] == 'system']
            
            # Get student language info
            student_info = self.db_service.get_student_info(student_id)
            language_info = {}
            curriculum_info = {}
            
            if student_info:
                language_info = {
                    "study_language": student_info['study_language'].value,
                    "nationality": student_info['nationality'],
                    "grade": student_info['grade']
                }
                
                # Add JSON-based curriculum availability info
                if self.pgvector:
                    curriculum = self.pgvector.get_curriculum_structure(
                        student_info['grade'], student_info['is_arabic']
                    )
                    curriculum_info = {
                        "curriculum_available": curriculum is not None,
                        "curriculum_source": "JSON file" if curriculum else "None",
                        "available_subjects": self.pgvector.get_subjects_by_grade_and_language(
                            student_info['grade'], student_info['is_arabic']
                        ),
                        "available_units": len(curriculum.get('units', [])) if curriculum else 0
                    }
            
            return {
                "student_id": student_id,
                "total_messages": len(history),
                "user_messages": len(user_messages),
                "assistant_messages": len(assistant_messages),
                "system_messages": len(system_messages),
                "conversation_active": len(history) > 0,
                **language_info,
                **curriculum_info
            }
        except Exception as e:
            logger.error(f"Error getting enhanced agent stats: {e}")
            return {
                "student_id": student_id,
                "error": str(e)
            }

    def get_curriculum_structure_info(self, student_id: str, subject: str = "Science") -> Dict:
        """Get detailed curriculum structure information from JSON data"""
        if not self.pgvector:
            return {"error": "PGVector service not available"}
            
        try:
            student_info = self.db_service.get_student_info(student_id)
            if not student_info:
                return {"error": "Student not found"}
            
            curriculum = self.pgvector.get_curriculum_structure(
                student_info['grade'], student_info['is_arabic'], subject
            )
            
            if not curriculum:
                return {
                    "error": f"No curriculum structure found",
                    "grade": student_info['grade'],
                    "language": "Arabic" if student_info['is_arabic'] else "English",
                    "subject": subject
                }
            
            # Extract detailed structure info
            units_info = []
            if 'units' in curriculum:
                for unit in curriculum['units']:
                    unit_info = {
                        "number": unit.get('number'),
                        "name": unit.get('name'),
                        "description": unit.get('description', ''),
                        "concepts_count": len(unit.get('concepts', [])),
                        "concepts": []
                    }
                    
                    for concept in unit.get('concepts', []):
                        concept_info = {
                            "number": concept.get('number'),
                            "name": concept.get('name'),
                            "description": concept.get('description', ''),
                            "lessons_count": len(concept.get('lessons', []))
                        }
                        unit_info["concepts"].append(concept_info)
                    
                    units_info.append(unit_info)
            
            return {
                "student_info": {
                    "grade": student_info['grade'],
                    "language": "Arabic" if student_info['is_arabic'] else "English",
                    "nationality": student_info['nationality']
                },
                "curriculum_title": curriculum.get('title', ''),
                "total_units": len(units_info),
                "units": units_info,
                "source": "JSON-based curriculum file"
            }
            
        except Exception as e:
            logger.error(f"Error getting curriculum structure info: {e}")
            return {"error": str(e)}

    def close(self):
        """Close database connection pools"""
        if self.pool_handler:
            try:
                self.pool_handler.close_all()
            except Exception as e:
                logger.error(f"Error closing connection pools: {e}")

    def __del__(self):
        """Destructor to ensure connection pools are closed"""
        self.close()