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
 

logger = logging.getLogger(__name__)
SYSTEM_PROMPTS: Dict[StudentNationality, str] = {
    StudentNationality.EGYPTIAN: """
إنت مُدرّس لطفل في ابتدائي اسمه {student_name}.  
رد باللهجة المصريّة الطبيعيّة كأنك بتكلم {student_name} قصادك.  
خلي الكلام بسيط، واضح، وقريب من ودنه.  
الجمل قصيرة ومترابطة، مش مقطّعة.  
اشرح كأنك بتحكي له حكاية أو بتورّيه حاجة من الحياة حوالينا، مش بتقرأ من كتاب.  

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

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

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

    StudentNationality.SAUDI: """
إنت معلّم لطفل في ابتدائي اسمه {student_name}.  
رد باللهجة السعوديّة الطبيعيّة، كأنك تشرح له قدّامك.  
خل الشرح واضح وسهل، لكن لا يكون ناشف.  
اشرح كأنك تسولف معه وتشبّه بأشياء من حياته اليومية.  

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

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

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








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

    def __init__(self, use_pgvector: bool = True):
        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

        # Use database for conversation memory
        self.db_service = ChatDatabaseService()
        self.pgvector = PGVectorService() if use_pgvector else None

    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"""
        return self.db_service.get_chat_history(student_id)

    def add_message_to_history(self, student_id: str, message: str, role: str = "user"):
        """Add message to database"""
        self.db_service.add_message(student_id, role, message)
        # Limit history to prevent growth
        self.db_service.limit_history(student_id, max_messages=38)

    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 []
        
        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']
        )

    def generate_response(
        self,
        user_message: str,
        student_id: str,
        subject: str = "Science",
        model: str = Models.chat,
        temperature: float = 1.0,
        top_k: int = 3
    ) -> str:
        """Generate AI response using database memory with subject filtering"""
        if not self.is_available():
            raise HTTPException(status_code=500, detail="Agent service not available")

        try:
            # Get complete student information from database
            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 student first name only
            full_name = student_info.get('student_name', 'الطالب')
            student_name = full_name.split()[0] if full_name else "الطالب"
            
            # Print student information
            print("----------------- Student Info Retrieved -----------------")
            print(f"Student ID: {student_id}")
            for key, value in student_info.items():
                print(f"{key.capitalize()}: {value}")
            print("---------------------------------------------------------")

            logger.info(f"Retrieved student info from DB: {student_info} for student: {student_id}")
            
            # Convert nationality string to StudentNationality enum
            nationality_lower = student_info['nationality'].lower().strip()
            nationality_mapping = {
                'egyptian': StudentNationality.EGYPTIAN,
                'saudi': StudentNationality.SAUDI
            }
            
            if nationality_lower in nationality_mapping:
                nationality = nationality_mapping[nationality_lower]
                logger.info(f"Successfully mapped '{student_info['nationality']}' to {nationality}")
            else:
                logger.warning(f"Unknown nationality '{student_info['nationality']}' for student {student_id}, defaulting to EGYPTIAN")
                nationality = StudentNationality.EGYPTIAN

            # Add user message to database
            self.add_message_to_history(student_id, user_message, "user")

            # Get conversation history from database
            conversation_history = self.get_conversation_history(student_id)

            # Create subject-specific system prompt with first name only
            base_system_prompt = SYSTEM_PROMPTS.get(nationality, SYSTEM_PROMPTS[StudentNationality.EGYPTIAN])
            formatted_base_prompt = base_system_prompt.format(student_name=student_name)
            subject_specific_prompt = f"{formatted_base_prompt}\n\nإنت بتدرّس مادة {subject} للطفل {student_name}."
            
            logger.info(f"Using nationality: {nationality}, subject: {subject}, and student name: {student_name} for student: {student_id}")

            # Prepare messages
            messages = []
            
            # Check if system message exists
            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})
                # Add system message to database
                self.add_message_to_history(student_id, subject_specific_prompt, "system")
            
            # Add conversation history
            messages.extend(conversation_history)

            # Enhanced pgvector enrichment with filtering
            if self.pgvector:
                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
                    )

                    if neighbors:
                        print("\n----------------- Retrieval Results -----------------")
                        context_message = f"معلومات من المنهج لمادة {subject} للصف {student_info['grade']} للطالب {student_name}:\n"
                        for i, n in enumerate(neighbors, 1):
                            unit_info = f" - الوحدة: {n['unit']}" if n['unit'] else ""
                            concept_info = f" - المفهوم: {n['concept']}" if n['concept'] else ""
                            lesson_info = f" - الدرس: {n['lesson']}" if n['lesson'] else ""
                            
                            context_message += f"\n{i}. {unit_info}{concept_info}{lesson_info}\n"
                            context_message += f"المحتوى: {n['chunk_text'][:200]}...\n"
                            context_message += f"(درجة التشابه: {n['distance']:.3f})\n"

                            print(f"Result {i}:")
                            print(f"  Unit: {n['unit']}")
                            print(f"  Concept: {n['concept']}")
                            print(f"  Lesson: {n['lesson']}")
                            print(f"  Chunk Text: {n['chunk_text']}...")
                            print(f"  Distance: {n['distance']:.3f}")
                            print("-" * 20)
                        print("-----------------------------------------------------")
                        
                        messages.append({"role": "system", "content": context_message})
                        logger.info(f"Added {len(neighbors)} filtered knowledge base results for subject: {subject}")
                    else:
                        print("\n----------------- Retrieval Results -----------------")
                        print(f"No relevant content found for subject: {subject}, grade: {student_info['grade']}, Arabic: {student_info['is_arabic']}")
                        print("-----------------------------------------------------")
                        logger.info(f"No relevant content found for subject: {subject}, grade: {student_info['grade']}, Arabic: {student_info['is_arabic']}")
                        
                except Exception as e:
                    logger.warning(f"Error using pgvector with filtering: {e}")

            # Generate AI 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")

            # Add AI response to database
            self.add_message_to_history(student_id, ai_response, "assistant")
            
            return ai_response

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

    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
            student_name = student_info.get('student_name', 'الطالب')
            
            # 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)
            subject_specific_prompt = f"{formatted_base_prompt}\n\nإنت بتدرّس مادة {subject} للطفل {student_name}."
            
            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 close(self):
        """Close database connections"""
        if self.db_service:
            self.db_service.close()
        if self.pgvector:
            self.pgvector.close()


# ----------------- Test -----------------
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    agent = AgentService(use_pgvector=True)

    if agent.is_available():
        try:
            # Test with chemistry (default)
            reply = agent.generate_response(
                "هو يعني إيه ذَرّة؟", 
                student_id="student_001",
                subject="chemistry"
            )
            print("AI (Chemistry):", reply)
            
            # Test with math
            reply = agent.generate_response(
                "إيه هو الجمع؟", 
                student_id="student_001",
                subject="math"
            )
            print("AI (Math):", reply)
            
        except Exception as e:
            print(f"Test failed: {e}")
        finally:
            agent.close()
    else:
        print("Agent service not available. Check OPENAI_API_KEY.")