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
from services.pedagogy_service import PedagogyService
from services.connection_pool import ConnectionPool

logger = logging.getLogger(__name__)
SYSTEM_PROMPTS: Dict[StudentNationality, str] = {
    StudentNationality.EGYPTIAN: """
إنت مُدرّس لطفل في ابتدائي اسمه {student_name} في الصف {grade}.  

فقط لو الطِّفل سأل عن هويتك بصراحة وواضح (مثل "إنت مين؟"، "عرّفني بنفسك"، "إنت بتعمل إيه هنا؟")،  
رُد بالنصّ الثابت ده:  
"أنا عَنان مؤسِّس شارع العلوم، وإنت هنا على مَنَصّة Science Street Lab،  
وأنا هنا عشان أساعدك تتعلَّم أي حاجة عايز تتعلَّمها في العلوم."  

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

أما لو سأل عن أي حاجة في العلوم أو المنهج، اشرح له بالطريقة التعليمية المناسبة.

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

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

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

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

    StudentNationality.SAUDI: """
إنت مُعلّم لطفل في ابتدائي اسمه {student_name} في الصف {grade}.  

فقط لو الطفل سأل عن هويتك بصراحة وواضح (مثل "إنت مين؟"، "عرّفني بنفسك"، "إنت وش تسوي هنا؟")،  
رُد بالنص الثابت هذا:  
"أنا عَنان مؤسِّس شارع العلوم، وإنت هنا على مَنَصّة Science Street Lab،  
وأنا هنا عشان أساعدك تتعلَّم أي حاجة عايز تتعلَّمها في العلوم."  

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

أما لو سأل عن أي شيء في العلوم أو المنهج، اشرح له بالطريقة التعليمية المناسبة.

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

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

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

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

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

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



class AgentService:
    """Service class for handling AI agent conversations using database memory"""

    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)
        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 get_available_subjects(self, student_id: str) -> List[str]:
        """Get available subjects for the student based on their grade and language"""
        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 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:
        """Generate AI response using database memory with optional retrieval based on question type"""
        if not self.is_available():
            raise HTTPException(status_code=500, detail="Agent service not available")

        try:
            # Get student info
            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 first name
            full_name = student_info.get('student_name', 'الطالب')
            student_name = full_name.split()[0] if full_name else "الطالب"

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

            # Format system prompt
            base_system_prompt = SYSTEM_PROMPTS.get(nationality, SYSTEM_PROMPTS[StudentNationality.EGYPTIAN])
            formatted_base_prompt = base_system_prompt.format(
                student_name=student_name,
                grade=student_info['grade']
            )
            subject_specific_prompt = f"{formatted_base_prompt}\n\nإنت بتدرّس مادة {subject} للطفل {student_name}."

            # Add Socratic instructions if any
            socratic_instructions = self.pedagogy_service.get_socratic_instructions(
                student_info['grade'], student_info['nationality']
            )
            if socratic_instructions:
                subject_specific_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": subject_specific_prompt})
                self.add_message_to_history(student_id, subject_specific_prompt, "system")
            messages.extend(conversation_history)

            # ----------------- DYNAMIC RETRIEVAL DECISION -----------------
            # Ask model to classify if retrieval needed
            try:
                classification_prompt = f"""
                صنف السؤال التالي: هل يحتاج لإجابة تعتمد على المنهج الدراسي أو محتوى المادة العلمية المتخصصة؟
                رد فقط بـ "YES" لو يحتاج retrieval من المحتوى التعليمي، و "NO" لو مجرد سؤال عام أو عن الشخصية أو محادثة عادية.
                
                أمثلة على أسئلة تحتاج "YES":
                - ما هو التمثيل الضوئي؟
                - اشرح لي الجهاز الهضمي
                - كيف تتكون الأمطار؟
                
                أمثلة على أسئلة تحتاج "NO":
                - إنت مين؟
                - إزيك؟ 
                - نت عارف انا مين؟
                - شكرا ليك
                
                السؤال: "{user_message}"
                """
                classification_response = self.client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[{"role": "user", "content": classification_prompt}],
                    temperature=0
                )
                classification_answer = classification_response.choices[0].message.content.strip().upper()
                need_retrieval = classification_answer == "YES"
                logger.info(f"Classification for '{user_message}': {classification_answer} (need_retrieval: {need_retrieval})")
            except Exception as e:
                logger.warning(f"Error in classification, defaulting to no retrieval: {e}")
                need_retrieval = False

            # ----------------- RETRIEVAL (if needed) -----------------
            if self.pgvector and need_retrieval:
                try:
                    query_embedding = self.openai_service.generate_embedding(user_message)
                    neighbors = self.pgvector.search_filtered_nearest(
                        query_embedding=query_embedding,
                        grade=student_info['grade'],
                        subject=subject,
                        is_arabic=student_info['is_arabic'],
                        limit=top_k
                    )

                    relevant_neighbors = [n for n in neighbors if n['distance'] < 1.3] if neighbors else []

                    if relevant_neighbors:
                        context_message = f"معلومات من المنهج لمادة {subject} للصف {student_info['grade']}:\n\n"
                        for n in relevant_neighbors:
                            unit_info = f"الوحدة: {n['unit']}" if n.get('unit') else ""
                            concept_info = f"المفهوم: {n['concept']}" if n.get('concept') else ""
                            lesson_info = f"الدرس: {n['lesson']}" if n.get('lesson') else ""
                            context_header = " - ".join(filter(None, [unit_info, concept_info, lesson_info]))
                            if context_header:
                                context_message += f"**{context_header}**\n"
                            context_message += f"{n['chunk_text']}\n\n---\n\n"
                        context_message += "استخدم هذه المعلومات لتقديم شرح دقيق ومناسب للطفل."
                        messages.append({"role": "system", "content": context_message})

                except Exception as e:
                    logger.warning(f"Error using pgvector: {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")
            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"""
        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")
            
            return self.pgvector.search_filtered_nearest(
                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 search_similar: {e}")
            raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")

    def update_student_subject_context(self, student_id: str, subject: str):
        """Update the system message for a new subject"""
        try:
            student_info = self.db_service.get_student_info(student_id)
            if not student_info:
                return False
            
            # Extract student name
            full_name = student_info.get('student_name', 'الطالب')
            student_name = full_name.split()[0] if full_name else "الطالب"
            
            # Clear existing history to reset context
            self.db_service.clear_history(student_id)
            
            # Set new system message with subject and student name
            nationality_lower = student_info['nationality'].lower().strip()
            nationality_mapping = {
                'egyptian': StudentNationality.EGYPTIAN,
                'saudi': StudentNationality.SAUDI
            }
            
            nationality = nationality_mapping.get(nationality_lower, StudentNationality.EGYPTIAN)
            base_system_prompt = SYSTEM_PROMPTS.get(nationality, SYSTEM_PROMPTS[StudentNationality.EGYPTIAN])
            # Format the prompt with student name
            formatted_base_prompt = base_system_prompt.format(
                student_name=student_name,
                grade=student_info['grade']
            )

            subject_specific_prompt = f"{formatted_base_prompt}\n\nإنت بتدرّس مادة {subject} للطفل {student_name}."
            
            # Add Socratic questioning instructions if applicable
            socratic_instructions = self.pedagogy_service.get_socratic_instructions(
                student_info['grade'], 
                student_info['nationality']
            )
            if socratic_instructions:
                subject_specific_prompt += f"\n\n{socratic_instructions}"
                logger.info(f"Added Socratic instructions when updating subject context for grade {student_info['grade']}")
            
            self.add_message_to_history(student_id, subject_specific_prompt, "system")
            
            logger.info(f"Updated subject context to {subject} for student {student_id} ({student_name})")
            return True
            
        except Exception as e:
            logger.error(f"Error updating subject context: {e}")
            return False

    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 import_conversation(self, messages: List[Dict[str, str]], student_id: str):
        """Import conversation history for a student"""
        try:
            # Clear existing history first
            self.db_service.clear_history(student_id)
            
            # Import messages in order
            for message in messages:
                role = message.get("role", "user")
                content = message.get("content", "")
                if content:
                    self.add_message_to_history(student_id, content, role)
        except Exception as e:
            logger.error(f"Error importing conversation for {student_id}: {e}")
            raise

    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"""
        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']
            
            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
            }
        except Exception as e:
            logger.error(f"Error getting agent stats: {e}")
            return {
                "student_id": student_id,
                "error": str(e)
            }

    def set_system_prompt(self, prompt: str) -> Dict[str, str]:
        """Set a new system prompt (this is a placeholder - actual implementation would depend on requirements)"""
        # This method would need to be implemented based on your specific requirements
        # for how system prompts should be managed globally vs per student
        return {
            "status": "success",
            "message": "System prompt updated"
        }

    @property
    def system_prompt(self) -> str:
        """Get the current system prompt"""
        # Return a default system prompt - this could be made more sophisticated
        return "Default system prompt for educational AI assistant"

    def debug_retrieval_pipeline(self, student_id: str, query: str):
        """Debug function to trace the retrieval pipeline"""
        print("=== RETRIEVAL DEBUG PIPELINE ===")
        
        try:
            # 1. Check student info
            student_info = self.db_service.get_student_info(student_id)
            print(f"1. Student Info: {student_info}")
            
            if not student_info:
                print("❌ No student info found!")
                return
            
            # 2. Test embedding generation
            print(f"2. Testing embedding generation for query: '{query}'")
            query_embedding = self.openai_service.generate_embedding(query)
            print(f"✅ Generated embedding (length: {len(query_embedding)})")
            
            # 3. Test vector search
            print(f"3. Testing vector search...")
            if self.pgvector:
                neighbors = self.pgvector.search_filtered_nearest(
                    query_embedding=query_embedding,
                    grade=student_info['grade'],
                    subject="Science",
                    is_arabic=student_info['is_arabic'],
                    limit=3
                )
                
                print(f"✅ Found {len(neighbors)} neighbors:")
                for i, neighbor in enumerate(neighbors):
                    print(f"   {i+1}. Distance: {neighbor['distance']:.3f}")
                    print(f"      Unit: {neighbor.get('unit', 'N/A')}")
                    print(f"      Concept: {neighbor.get('concept', 'N/A')}")
                    print(f"      Text preview: {neighbor['chunk_text'][:100]}...")
                    print()
            else:
                print("❌ PGVector not available")
            
            # 4. Test full response generation
            print(f"4. Testing full response generation...")
            response = self.generate_response(
                user_message=query,
                student_id=student_id,
                subject="Science"
            )
            print(f"✅ Generated response (first 200 chars): {response[:200]}...")
            
            print("=== DEBUG COMPLETE ===")
            
        except Exception as e:
            print(f"❌ Debug pipeline failed at step: {e}")
            import traceback
            traceback.print_exc()

    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()