# src/agent.py
from llama_cpp import Llama
from src.memory import MemoryManager
import os
import logging
from src.utils import extract_and_summarize  # Import extract_and_summarize
from src.prompts import Prompts  # Import system prompts

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class Agent:
    def __init__(self, llm: Llama, db_path: str, system_prompt: str = "", max_tokens: int = 500, temperature: float = 0.7, top_p: float = 0.95):
        self.llm = llm
        self.memory = MemoryManager(db_path)
        self.prompts = {
            "general": Prompts.GENERAL_SYSTEM_PROMPT,
            "whole_document": Prompts.WHOLE_DOCUMENT_SYSTEM_PROMPT,
            "query_response": Prompts.QUERY_RESPONSE_SYSTEM_PROMPT,
            "enhancement": Prompts.ENHANCEMENT_SYSTEM_PROMPT
        }
        self.max_tokens = max_tokens  # Default value
        self.temperature = temperature  # Default value
        self.top_p = top_p  # Default value

    def process_query(self, user_id: str, query: str) -> str:
        # Normalize the query to lowercase
        query = query.lower()

        # Check if the query is specific (e.g., "chronique #441")
        if "chronique #" in query:
            section_number = query.split("chronique #")[1].split()[0]
            section_description = self.get_section_description(f"chronique #{section_number}")
            if section_description:
                response = self.generate_specific_response(query, section_description)
                self.memory.add_user_interaction(user_id, query, response)
                return response

        if "flash info fl-" in query:
            section_number = query.split("flash info fl-")[1].split()[0]
            section_description = self.get_section_description(f"flash info fl-{section_number}")
            if section_description:
                response = self.generate_specific_response(query, section_description)
                self.memory.add_user_interaction(user_id, query, response)
                return response

        if "chronique-faq #" in query:
            section_number = query.split("chronique-faq #")[1].split()[0]
            section_description = self.get_section_description(f"chronique-faq #{section_number}")
            if section_description:
                response = self.generate_specific_response(query, section_description)
                self.memory.add_user_interaction(user_id, query, response)
                return response

        # For general queries, use the existing multi-layer processing
        initial_response = extract_and_summarize(query, self.memory, self.llm, self.dynamic_query_response_prompt(query), max_tokens=self.max_tokens, temperature=self.temperature, top_p=self.top_p)

        # Evaluate the initial response
        if not self.evaluate_response(initial_response, query):
            # Retrieve additional data
            additional_data = self.retrieve_additional_data(query, initial_response)
            # Combine initial and additional data
            combined_context = f"{initial_response}\n{additional_data}"
            # Truncate the combined context to fit within the model's context window
            max_context_length = 30000  # Adjust this based on your LLM's token limit
            if len(combined_context) > max_context_length:
                combined_context = combined_context[:max_context_length]
                logging.info(f"Truncated combined context to {max_context_length} characters.")

            # Generate final response
            initial_response = self.llm.create_chat_completion(
                messages=[{"role": "user", "content": f"Context: {combined_context}\nQuestion: {query}"}],
                max_tokens=self.max_tokens,
                temperature=self.temperature,
                top_p=self.top_p
            )['choices'][0]['message']['content']

        refined_response = self.multi_layer_processing(query, initial_response)
        self.memory.add_user_interaction(user_id, query, refined_response)
        return refined_response

    def generate_specific_response(self, query: str, section_description: str) -> str:
        # Log the section description
        logging.info(f"Found section: {query}")
        logging.info(f"Section description for {query}: {section_description}")

        # Format the context with a placeholder pre-set
        context = f"Section description for {query}: {section_description}"

        # Generate a response using the specific section description
        response = self.llm.create_chat_completion(
            messages=[{"role": "user", "content": f"Context: {context}\nQuestion: {query}"}],
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p
        )['choices'][0]['message']['content']

        # Log the generated response
        logging.info(f"Generated specific response: {response}")

        return response

    def multi_layer_processing(self, query: str, initial_response: str) -> str:
        # First layer: High-level summary of the entire document
        high_level_summary = self.generate_high_level_summary(query)

        # Second layer: Refine the response using the high-level summary
        refined_response = self.refine_response(query, initial_response, high_level_summary)

        return refined_response

    def generate_high_level_summary(self, query: str) -> str:
        # Retrieve all memories from the database
        all_memories = self.memory._get_all_memories()

        # Count the number of each type of section
        chronique_count = self.count_chroniques()
        flash_info_count = self.count_flash_infos()
        chronique_faq_count = self.count_chronique_faqs()

        # Combine all descriptions into a single context
        full_context = " ".join([memory['description'] for memory, _, _ in all_memories])

        # Truncate the context if it exceeds the token limit
        max_context_length = 500  # Adjust this based on your LLM's token limit
        if len(full_context) > max_context_length:
            full_context = full_context[:max_context_length]
            logging.info(f"Truncated full context to {max_context_length} characters.")

        # Generate a high-level summary using the LLM
        high_level_summary = self.llm.create_chat_completion(
            messages=[{"role": "user", "content": f"Context: {full_context}\nQuestion: {query}"}],
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p
        )['choices'][0]['message']['content']

        # Explicitly include the number of sections in the summary
        high_level_summary += f"\nD'après les données disponibles, Michel Thomas a publié {chronique_count} chroniques, {flash_info_count} flash infos, et {chronique_faq_count} chronique-faqs."

        logging.info(f"Generated high-level summary: {high_level_summary}")

        return high_level_summary

    def refine_response(self, query: str, initial_response: str, high_level_summary: str) -> str:
        # Combine the initial response and the high-level summary
        combined_context = f"Initial Response: {initial_response}\nHigh-Level Summary: {high_level_summary}"

        # Truncate the combined context to fit within the model's context window
        max_context_length = 500  # Adjust this based on your LLM's token limit
        if len(combined_context) > max_context_length:
            combined_context = combined_context[:max_context_length]
            logging.info(f"Truncated combined context to {max_context_length} characters.")

        # Generate a refined response using the LLM
        refined_response = self.llm.create_chat_completion(
            messages=[{"role": "user", "content": f"Context: {combined_context}\nQuestion: {query}"}],
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p
        )['choices'][0]['message']['content']

        # Enhance the response to include more details about the specific section
        refined_response = self.enhance_response_with_details(query, refined_response)

        logging.info(f"Generated refined response: {refined_response}")

        return refined_response

    def enhance_response_with_details(self, query: str, refined_response: str) -> str:
        if "chronique #" in query:
            section_number = query.split("chronique #")[1].split()[0]
            section_description = self.get_section_description(f"chronique #{section_number}")
        elif "flash info fl-" in query:
            section_number = query.split("flash info fl-")[1].split()[0]
            section_description = self.get_section_description(f"flash info fl-{section_number}")
        elif "chronique-faq #" in query:
            section_number = query.split("chronique-faq #")[1].split()[0]
            section_description = self.get_section_description(f"chronique-faq #{section_number}")
        else:
            section_description = ""

        if section_description:
            refined_response += f"\n\nVoici plus de détails sur la section demandée:\n{section_description}"

        return refined_response

    def get_section_description(self, section_name: str) -> str:
        # Normalize the section name to lowercase
        section_name = section_name.lower()

        # Retrieve the specific section from the database
        self.memory.cursor.execute("SELECT description FROM semantic_memory WHERE concept LIKE ?", (f"{section_name}%",))
        result = self.memory.cursor.fetchone()
        if result:
            logging.info(f"Found section: {section_name}")
            return result[0]
        else:
            logging.warning(f"Section not found: {section_name}")
            return ""

    def enhance_response(self, query: str, initial_response: str) -> str:
        # Ask the LLM to enhance the response
        try:
            enhanced_query = f"{initial_response} Comment pourriez-vous améliorer cette réponse pour qu'elle soit entièrement conforme à la requête de l'utilisateur ? {query}. Notez que vous êtes l'assistant IA de Michel Thomas, le consultant qui a écrit tout ce contenu sur son site web, y compris toutes les Chroniques, Flash Infos et Chronique-FAQ sur https://mtc-qc.ca. Votre tâche est de formuler une réponse claire et concise pour les utilisateurs du site web. - Utilisez uniquement le contenu de la base de données. Ne jamais utiliser des connaissances externes ou des suppositions. Si le contenu de la base de données ne fournit pas suffisamment d'informations, indiquez clairement que vous ne disposez pas de suffisamment d'informations pour répondre à la question. - Ne jamais ecrire des charactere semblable a  (c)  car cela reprensente le charactere  é . "
            enhanced_response = self.llm.create_chat_completion(
                messages=[{"role": "user", "content": enhanced_query}],
                max_tokens=self.max_tokens,
                temperature=self.temperature,
                top_p=self.top_p
            )['choices'][0]['message']['content']
            logging.info(f"Enhanced response: {enhanced_response}")
        except Exception as e:
            enhanced_response = f"Erreur lors de l'amélioration de la réponse: {e}"
            logging.error(f"Erreur lors de l'amélioration de la réponse: {e}")

        return enhanced_response

    def evaluate_response(self, response: str, query: str) -> bool:
        # Evaluate the response to determine if it is sufficient
        # Example: Check if the response contains all key terms from the query
        key_terms = set(query.split())
        response_terms = set(response.split())
        return key_terms.issubset(response_terms)

    def retrieve_additional_data(self, query: str, initial_response: str) -> str:
        # Retrieve additional relevant data from the database
        key_terms = set(query.split()) | set(initial_response.split())
        relevant_memories = self.memory.retrieve_relevant_memories(" ".join(key_terms), limit=10)
        additional_data = " ".join([memory['description'] for memory in relevant_memories])
        return additional_data

    def dynamic_query_response_prompt(self, query: str) -> str:
        return f"""
        Vous êtes l'assistant intelligent de Michel Thomas. Votre tâche est de répondre à la requête de l'utilisateur en utilisant uniquement le contexte fourni dans la base de données. La requête de l'utilisateur est: "{query}". Assurez-vous que vos réponses sont claires, précises et directement liées à la requête de l'utilisateur. Si possible, incluez des exemples concrets pour illustrer vos points.

        - Utilisez uniquement le contenu de la base de données pour générer la réponse.
        - Ne jamais utiliser des connaissances externes ou des suppositions.
        - Résumez le contenu de manière concise et claire.
        - Expliquez les concepts de manière détaillée et accessible, en utilisant des exemples concrets et des analogies.
        - Analysez le contenu en profondeur, identifiez les thèmes principaux et les arguments clés.
        - Critiquez le contenu de manière constructive, en identifiant les points forts et les points faibles.
        - Générez des questions pertinentes pour encourager une réflexion plus approfondie.
        - Fournissez des exemples concrets pour illustrer les concepts.
        - Formulez des hypothèses basées sur les informations disponibles.
        - Tirez des conclusions bien fondées et soutenues par des preuves.
        - Formulez des recommandations pratiques et applicables.
        - Générez des éléments d'action spécifiques, mesurables, réalisables, pertinents et temporellement définis (SMART).
        - Créez un plan détaillé pour organiser les idées principales et les sous-thèmes.
        - Rédigez une introduction engageante pour captiver l'attention du lecteur.
        - Rédigez une conclusion forte pour résumer les points clés et laisser une impression durable.
        - Paraphrasez le contenu de manière claire et concise, en préservant le sens original.
        - Simplifiez le contenu pour le rendre plus accessible à un public plus large.
        - Développez le contenu en ajoutant plus de détails et d'informations.
        - Comparez deux concepts ou sujets en identifiant les similitudes et les différences.
        - Contrastez deux concepts ou sujets en mettant en évidence les différences significatives.
        - Créez des analogies pertinentes pour clarifier les concepts.
        - Créez des métaphores pertinentes pour illustrer les concepts de manière créative.
        - Si le contenu de la base de données ne fournit pas suffisamment d'informations, indiquez clairement que vous ne disposez pas de suffisamment d'informations pour répondre à la question.
        """

    def dynamic_whole_document_prompt(self, query: str, chronique_count: int, flash_info_count: int, chronique_faq_count: int) -> str:
        return f"""
        Vous êtes l'assistant intelligent de Michel Thomas, consultant qui a écrit tout ce contenu sur son site web, y compris toutes les Chroniques, Flash Infos et Chronique-FAQ sur https://mtc-qc.ca. Votre tâche est de comprendre l'ensemble du document et de générer un résumé de haut niveau ou un contexte qui peut être utilisé pour répondre à la requête de l'utilisateur. La requête de l'utilisateur est: "{query}". Concentrez-vous sur les points les plus pertinents et importants. Incluez le nombre total de Chroniques ({chronique_count}), Flash Infos ({flash_info_count}), et Chronique-FAQ ({chronique_faq_count}) publiées par Michel Thomas.

        - Utilisez uniquement le contenu de la base de données pour générer le résumé.
        - Ne jamais utiliser des connaissances externes ou des suppositions.
        - Résumez le contenu de manière concise et claire.
        - Expliquez les concepts de manière détaillée et accessible, en utilisant des exemples concrets et des analogies.
        - Analysez le contenu en profondeur, identifiez les thèmes principaux et les arguments clés.
        - Critiquez le contenu de manière constructive, en identifiant les points forts et les points faibles.
        - Générez des questions pertinentes pour encourager une réflexion plus approfondie.
        - Fournissez des exemples concrets pour illustrer les concepts.
        - Formulez des hypothèses basées sur les informations disponibles.
        - Tirez des conclusions bien fondées et soutenues par des preuves.
        - Formulez des recommandations pratiques et applicables.
        - Générez des éléments d'action spécifiques, mesurables, réalisables, pertinents et temporellement définis (SMART).
        - Créez un plan détaillé pour organiser les idées principales et les sous-thèmes.
        - Rédigez une introduction engageante pour captiver l'attention du lecteur.
        - Rédigez une conclusion forte pour résumer les points clés et laisser une impression durable.
        - Paraphrasez le contenu de manière claire et concise, en préservant le sens original.
        - Simplifiez le contenu pour le rendre plus accessible à un public plus large.
        - Développez le contenu en ajoutant plus de détails et d'informations.
        - Comparez deux concepts ou sujets en identifiant les similitudes et les différences.
        - Contrastez deux concepts ou sujets en mettant en évidence les différences significatives.
        - Créez des analogies pertinentes pour clarifier les concepts.
        - Créez des métaphores pertinentes pour illustrer les concepts de manière créative.
        - Si le contenu de la base de données ne fournit pas suffisamment d'informations, indiquez clairement que vous ne disposez pas de suffisamment d'informations pour répondre à la question.
        """

    def dynamic_enhancement_prompt(self, query: str) -> str:
        return f"""
        Vous êtes l'assistant intelligent de Michel Thomas. Votre tâche est d'améliorer la réponse initiale en la rendant plus complète et plus conforme à la requête de l'utilisateur. La requête de l'utilisateur est: "{query}". Prenez en compte tous les détails pertinents et formulez votre réponse de manière concise et claire. Ajoutez des détails supplémentaires si nécessaire pour rendre la réponse plus informative.

        - Utilisez uniquement le contenu de la base de données pour améliorer la réponse.
        - Ne jamais utiliser des connaissances externes ou des suppositions.
        - Résumez le contenu de manière concise et claire en maximum 8 phrases de maximum 88 mots chacun.
        - Expliquez les concepts de manière détaillée et accessible, en utilisant des exemples concrets et des analogies.
        - Analysez le contenu en profondeur, identifiez les thèmes principaux et les arguments clés.
        - Critiquez le contenu de manière constructive, en identifiant les points forts et les points faibles.
        - Générez des questions pertinentes pour encourager une réflexion plus approfondie.
        - Fournissez des exemples concrets pour illustrer les concepts.
        - Formulez des hypothèses basées sur les informations disponibles.
        - Tirez des conclusions bien fondées et soutenues par des preuves.
        - Formulez des recommandations pratiques et applicables.
        - Générez des éléments d'action spécifiques, mesurables, réalisables, pertinents et temporellement définis (SMART).
        - Créez un plan détaillé pour organiser les idées principales et les sous-thèmes.
        - Rédigez une introduction engageante pour captiver l'attention du lecteur.
        - Rédigez une conclusion forte pour résumer les points clés et laisser une impression durable.
        - Paraphrasez le contenu de manière claire et concise, en préservant le sens original.
        - Simplifiez le contenu pour le rendre plus accessible à un public plus large.
        - Développez le contenu en ajoutant plus de détails et d'informations.
        - Comparez deux concepts ou sujets en identifiant les similitudes et les différences.
        - Contrastez deux concepts ou sujets en mettant en évidence les différences significatives.
        - Créez des analogies pertinentes pour clarifier les concepts.
        - Créez des métaphores pertinentes pour illustrer les concepts de manière créative.
        - Si le contenu de la base de données ne fournit pas suffisamment d'informations, indiquez clairement que vous ne disposez pas de suffisamment d'informations pour répondre à la question.
        """

    def count_chroniques(self) -> int:
        # Count the number of chroniques in the database
        self.memory.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE 'chronique #%'")
        count = self.memory.cursor.fetchone()[0]
        logging.info(f"Number of chroniques: {count}")
        return count

    def count_flash_infos(self) -> int:
        # Count the number of flash infos in the database
        self.memory.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE 'flash info fl-%'")
        count = self.memory.cursor.fetchone()[0]
        logging.info(f"Number of flash infos: {count}")
        return count

    def count_chronique_faqs(self) -> int:
        # Count the number of chronique-faqs in the database
        self.memory.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE 'chronique-faq #%'")
        count = self.memory.cursor.fetchone()[0]
        logging.info(f"Number of chronique-faqs: {count}")
        return count