import os
import sys
from typing import Dict
from fastapi import HTTPException
import concurrent.futures 
from concurrent.futures import ThreadPoolExecutor
from services.agent_helpers.agent_prompts import SYSTEM_PROMPTS
import logging
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
from schemas import StudentProfile
from core import StudentNationality, StudyLanguage, Models

logger = logging.getLogger(__name__)

class ResponseGenerator:
    """Handles AI response generation and conversation management"""

    def __init__(self, openai_service, db_service, pedagogy_service, query_handler, context_generator, agent_service):
        self.openai_service = openai_service
        self.db_service = db_service
        self.pedagogy_service = pedagogy_service
        self.query_handler = query_handler
        self.context_generator = context_generator
        self.agent_service = agent_service

    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 prepare_system_prompt(self, student_info: StudentProfile) -> str:
        """Prepare system prompt based on student information"""
        # Dot notation
        student_name = student_info.student_name.split()[0]
        study_language = student_info.study_language
        nationality = student_info.nationality

        # Get appropriate system prompt
        prompt_key = (nationality, study_language)
        base_system_prompt = SYSTEM_PROMPTS.get(prompt_key, 
            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
        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}"

        return formatted_base_prompt

    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 response generation with Parallel Speculative Execution.
        Classification and Vector Search run simultaneously.
        """
        if not self.openai_service.is_available():
            raise HTTPException(status_code=500, detail="Agent service not available")

        try:
            # 1. prepare (Sequential - fast DB lookups)
            # Fetch student info and history first as they are needed for inputs
            with ThreadPoolExecutor(max_workers=2) as pre_executor:
                future_student = pre_executor.submit(self.db_service.get_student_info, student_id)
                # Fetch history using the optimized 
                future_history = pre_executor.submit(self.get_conversation_history, student_id)
                
                student_info = future_student.result()
                conversation_history = future_history.result()

            if not student_info:
                raise HTTPException(status_code=404, detail=f"Student with ID {student_id} not found")
            
            
            # Save user message (Fire and forget - strictly speaking we should wait, 
            # but for speed we can do this without blocking the logic if you trust your DB)
            self.add_message_to_history(student_id, user_message, "user")

            # 2. speculative parallel execution 
            
            query_type = "specific_content" # Default fallback
            relevant_results = [] # Default empty
            
            with ThreadPoolExecutor(max_workers=2) as executor:
                
                # TASK A: THE BRAIN (Classification)
                # We pass the conversation_history we already fetched to save the DB call
                future_classification = executor.submit(
                    self.query_handler.classify_query_type,
                    user_message, 
                    student_info, 
                    student_id, 
                    conversation_history 
                )

                # TASK B: THE EYES (Vector Search)
                # We "Speculate" that the user MIGHT ask a content question.
                # We run the search immediately. If it turns out to be "general_chat", 
                # we just threw away 200ms of compute, but saved 1.5s of latency for real questions.
                future_search = executor.submit(
                    self.context_generator.search_enhanced_content,
                    user_message,
                    student_info,
                    subject,
                    top_k
                )

                # TASK C (Optional Check): Rule-based check for game help is fast, 
                # but we can do it inside the classification logic or separately.
                # Here we wait for the classification to finish.
                query_type = future_classification.result()
                
                logger.info(f"Query classified as: {query_type}")

                # 3. SYNCHRONIZATION & DECISION
                
                system_context_content = ""
                
                if query_type == "specific_content":
                    # The brain says: "This is science!"
                    # We check the eyes (search results). They are ALREADY READY.
                    relevant_results = future_search.result()
                    
                    if relevant_results:
                        system_context_content = self.context_generator.generate_enhanced_context(
                            relevant_results, student_info, query_type
                        )
                        logger.info(f"Using speculative search results: {len(relevant_results)} chunks")

                elif query_type == "game_help":
                    # Handle game help
                    game_context, user_query = self.query_handler.handle_game_help_query(user_message)
                    system_context_content = f"سياق اللعبة التعليمية اللي هتساعد الطفل فيها:\n{game_context}"
                    
                    # For game help, we also use the search results we started in Task B!
                    # (Assuming the game query might need knowledge)
                    relevant_results = future_search.result() 
                    if relevant_results:
                        enhanced_ctx = self.context_generator.generate_enhanced_context(
                            relevant_results, student_info, query_type
                        )
                        system_context_content += f"\n\nمحتوي المنهج اللي ليه علاقة بسؤال الطفل:\n{enhanced_ctx}"

                elif query_type == "general_chat":
                    # The brain says: "Just chatting."
                    # We IGNORE the search results from Task B. 
                    # (Task B might still be running or finished, we just don't use the data).
                    chat_context = self.query_handler.handle_general_chat_query(user_message, student_info)
                    system_context_content = f"سياق المحادثة العامة:\n{chat_context}"

                elif query_type == "overview":
                    overview_response = self.query_handler.handle_overview_query(student_info, subject)
                    system_context_content = f"المنهج الكامل من ملف JSON:\n{overview_response}"

                elif query_type == "navigation":
                    nav_response = self.query_handler.handle_navigation_query(user_message, student_info, subject)
                    system_context_content = f"تفاصيل الوحدة/المفهوم من JSON:\n{nav_response}"
                
                elif query_type == "ask_for_question":
                    # Special case, returns dict immediately
                    return {
                        "type": "mcq",
                        "data": self.agent_service.handle_ask_for_question(student_id)
                    }

            # 4. FINAL GENERATION

            # Prepare system prompt
            formatted_base_prompt = self.prepare_system_prompt(student_info)
            
            # Build messages
            messages = [{"role": "system", "content": formatted_base_prompt}]
            messages.extend(conversation_history)
            
            # Add the context we decided on (if any)
            if system_context_content:
                messages.append({"role": "system", "content": system_context_content})

            # Add user message
            messages.append({"role": "user", "content": user_message})

            # Call AI
            response = self.openai_service.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 response for student {student_id}: {e}")
            raise HTTPException(status_code=500, detail="Error generating response")