"""
SmartWebSearch.RAGTool
~~~~~~~~~~~~

This module implements the RAG tool.
"""

# Import the required modules
import numpy as np
from sentence_transformers import SentenceTransformer
from langchain_text_splitters import RecursiveCharacterTextSplitter
from typing import Any
from SmartWebSearch.Debugger import show_debug
from SmartWebSearch.Progress import Progress, _ProgressData
from SmartWebSearch.Progress import ProgressStatusSelector as pss
import datetime
import pickle
from pathlib import Path

class _KnowledgeBase:
    """
    A class for managing a knowledge base.
    """

    def __init__(self, knowledge_base: list[str], knowledge_vector: np.ndarray) -> None:
        """
        Initialize the KnowledgeBase object.

        Args:
            knowledge_base (list[str]): The knowledge base.
            knowledge_vector (np.ndarray): The knowledge vectors.

        Returns:
            None
        """
        self.knowledge_base: list[str] = knowledge_base
        self.knowledge_vector: np.ndarray = knowledge_vector

    def match_knowledge(self, embedding_model: SentenceTransformer, prompt: str, top_k: int = 10, threshold_score: float = 0.75):
        """
        Match the prompt with the knowledge base.

        Args:
            embedding_model (SentenceTransformer): The embedding model.
            prompt (str): The prompt to be matched.
            top_k (int) = 10: The number of top matches to return.
            threshold_score (float) = 0.75: The threshold score for a match.

        Returns:
            list[tuple[float, str]]: The top matches with their scores and the corresponding chunks.
        """

        # Encode the prompt into a vector (added 'query' as a prefix)
        prompt_vector: np.ndarray = embedding_model.encode(f"query: {prompt}")

        # Loop through the knowledge base and find the top matches
        matches: list[tuple[float, str]] = []
        for i, knowledge_vector in enumerate(self.knowledge_vector):
            score: float = np.dot(knowledge_vector, prompt_vector)
            matches.append((score, self.knowledge_base[i].strip()))

        # Sort the matches by score
        matches.sort(key = lambda knowledge: knowledge[0], reverse = True)

        # Return the top matches with a threshold score
        return [match for match in matches[:top_k] if match[0] > threshold_score]

class _KnowledgeBaseSet:
    """
    A class for managing a knowledge base set.
    """

    def __init__(self, knowledge_base_set: list[_KnowledgeBase]) -> None:
        """
        Initialize the KnowledgeBaseSet object.

        Args:
            knowledge_base_set (list[_KnowledgeBase]): The knowledge base set.

        Returns:
            None
        """
        self.knowledge_base_set: list[_KnowledgeBase] = knowledge_base_set

    def merge_knowledge(self, knowledge_base_set: "_KnowledgeBaseSet"):
        """
        Merge the knowledge base into the knowledge base set.

        Args:
            knowledge_base_set (_KnowledgeBaseSet): The knowledge base set to be merged.

        Returns:
            None
        """
        self.knowledge_base_set.extend(knowledge_base_set.knowledge_base_set)

    def match_knowledge(self, embedding_model: SentenceTransformer, prompt: str, top_k: int = 10, threshold_score: float = 0.75):
        """
        Match the prompt with the knowledge base set.

        Args:
            embedding_model (SentenceTransformer): The embedding model.
            prompt (str): The prompt to be matched.
            top_k (int) = 10: The number of top matches to return.
            threshold_score (float) = 0.75: The threshold score for a match.

        Returns:
            list[tuple[float, str]]: The top matches with their scores and the corresponding chunks.
        """

        # Match the prompt with the knowledge base set
        matches: list[tuple[float, str]] = []
        for knowledge_base in self.knowledge_base_set:
            matches.extend(knowledge_base.match_knowledge(embedding_model, prompt, top_k, threshold_score))

        # Sort the matches by score
        matches.sort(key = lambda knowledge: knowledge[0], reverse = True)

        # Return the top matches
        return [match for match in matches[:top_k + 5]]
    
    def save_knowledge(self, dir: str = "", filename: str = "knowledge_base") -> None:
        """
        Save the knowledge base set to a file.

        Args:
            dir (str) = "": The directory to save the knowledge base set.
            filename (str) = "knowledge_base": The filename to save the knowledge base set.

        Returns:
            None
        """

        # Save the knowledge base set to a file
        with open(Path(dir, f"{filename}.klb"), "wb") as file:
            pickle.dump(self, file)

    def __add__(self, knowledge_base_set: "_KnowledgeBaseSet") -> "_KnowledgeBaseSet":
        """
        Merge the knowledge base set into a new knowledge base set.

        Args:
            knowledge_base_set (_KnowledgeBaseSet): The knowledge base set to be merged.

        Returns:
            _KnowledgeBaseSet: The merged knowledge base set.
        """

        # Merge the knowledge base set into the knowledge base set
        return _KnowledgeBaseSet(self.knowledge_base_set + knowledge_base_set.knowledge_base_set)

# The RAGTool class
class RAGTool:
    """
    A class for RAG (Retrieval-Augmented Generation).
    """

    @staticmethod
    def load_knowledge(path: str = "knowledge_base.klb") -> _KnowledgeBaseSet:
        """
        Load the knowledge base set from a file.

        Args:
            path (str) = "knowledge_base.klb": The path to load the knowledge base set.

        Returns:
            _KnowledgeBaseSet: The knowledge base set.
        """

        # Load the knowledge base set from a file
        with open(path, "rb") as file:
            return pickle.load(file)

    @staticmethod
    def __build_knowledge_base(text_data: str, embedding_model: SentenceTransformer, text_splitter: RecursiveCharacterTextSplitter, progress: Progress) -> _KnowledgeBaseSet:
        """
        Build the knowledge base from the text data.

        Args:
            text_data (str): The text data to be used as knowledge.
            embedding_model (SentenceTransformer): The embedding model.
            text_splitter (RecursiveCharacterTextSplitter): The text splitter.
            progress (Progress): The progress object.

        Returns:
            _KnowledgeBaseSet: The knowledge base set.
        """

        # Split the text data into chunks
        chunks: list[str] = text_splitter.split_text(text_data)

        # Remove the double spaces and newlines
        for idx, chunk in enumerate(chunks):
            while '  ' in chunks[idx]:
                chunks[idx] = chunks[idx].replace('  ', ' ')
            while '\n ' in chunks[idx]:
                chunks[idx] = chunks[idx].replace('\n ', '\n')
            while '\n\n' in chunks[idx]:
                chunks[idx] = chunks[idx].replace('\n\n', '\n')

        # Remove the chunks with less than 100 characters
        chunks: list[str] = [chunk for chunk in chunks if len(chunk) > 100]

        # Remove the chunks includes enabling javascript
        chunks: list[str] = [chunk for chunk in chunks if 'enable' not in chunk.lower() or 'javascript' not in chunk.lower()]

        # Remove the chunks includes enabling cookies
        chunks: list[str] = [chunk for chunk in chunks if 'enable' not in chunk.lower() or 'cookie' not in chunk.lower()]

        # Remove the chunks includes verifying humans
        chunks: list[str] = [chunk for chunk in chunks if 'verify' not in chunk.lower() or 'human' not in chunk.lower()]

        # Add prefix for each chunk
        chunks: list[str] = [f"passage: {chunk}" for chunk in chunks]

        # Seperate the chunks into several chunk sets every 30 chunks
        chunk_sets: list[list[str]] = [chunks[i: i + 30] for i in range(0, len(chunks), 30)]

        # Set a timedelta for storing the the completion time for each knowledge base
        timedelta: datetime.timedelta = datetime.timedelta()
        
        # Encode the chunk sets into vector sets
        knowledge_vector_set: list[np.ndarray] = []

        show_debug(f"Creating knowledge base set...")
        for idx, chunk_set in enumerate(chunk_sets, start = 1):
            # Get the start time
            start_time: datetime.datetime = datetime.datetime.now()

            # Encode the chunk set into a vector set
            knowledge_vector_set.append(embedding_model.encode(chunk_set))

            # Get the end time
            end_time: datetime.datetime = datetime.datetime.now()

            # Calculate the timedelta
            timedelta: datetime.timedelta = end_time - start_time if end_time - start_time > timedelta else timedelta

            # Update the progress
            progress._update_progress(pss.KL_BASE_CREATING, f"Created knowledge base set {idx}/{len(chunk_sets)}.", {
                "current": idx,
                "total": len(chunk_sets),
                "eta": (timedelta) * (len(chunk_sets) - idx)
            }, idx / len(chunk_sets))

            show_debug(f"Created knowledge base set {idx}/{len(chunk_sets)} {f'(eta. {(timedelta) * (len(chunk_sets) - idx)}) ...' if len(chunk_sets) - idx > 0 else ''}")    
        
        # Update the progress
        progress._update_progress(pss.KL_BASE_CREATED, f"Knowledge base set created.", {
            "total": len(chunk_sets)
        })

        show_debug(f"Knowledge base set created.")

        # Create knowledge base objects for each set
        knowledge_base_set: list[_KnowledgeBase] = [_KnowledgeBase(chunk_set, knowledge_vector_set) for chunk_set, knowledge_vector_set in zip(chunk_sets, knowledge_vector_set)]

        # Create a knowledge base set object
        knowledge_base_set: _KnowledgeBaseSet = _KnowledgeBaseSet(knowledge_base_set)

        return knowledge_base_set

    def __init__(self, embedding_model_name: str = 'intfloat/multilingual-e5-base') -> None:
        """
        Initialize the RAGTool object.
        
        Args:
            embedding_model_name (str) = 'intfloat/multilingual-e5-base': The name of the embedding model.
        
        Returns:
            None
        """

        # Initialize the text splitter
        self.text_splitter: RecursiveCharacterTextSplitter = RecursiveCharacterTextSplitter(
            chunk_size = 400,
            chunk_overlap = 50
        )

        # Initialize the SentenceTransformer model
        self.embedding_model: SentenceTransformer = SentenceTransformer(embedding_model_name)

        # Initialize the Progress object
        self.progress: Progress = Progress()

    def build_knowledge(self, text_data: str) -> _KnowledgeBaseSet:
        """
        Build the knowledge base from the text data.

        Args:
            text_data (str): The text data to be used as knowledge.

        Returns:
            _KnowledgeBaseSet: The knowledge base set.
        """

        # Build the knowledge base
        knowledge_base_set: _KnowledgeBaseSet = self.__build_knowledge_base(text_data, self.embedding_model, self.text_splitter, self.progress)

        # Update the progress
        self.progress._update_progress(pss.COMPLETED, f"Knowledge base set created.", {
            "knowledge_base_set": knowledge_base_set
        })

        self.progress._update_progress(pss.IDLE)

        return knowledge_base_set

    def match_knowledge(self, knowledge_base: _KnowledgeBase | _KnowledgeBaseSet, prompt: str, top_k: int = 10, threshold_score: float = 0.75) -> list[tuple[float, str]]:
        """
        Match the prompt with the knowledge base.

        Args:
            knowledge_base (_KnowledgeBase | _KnowledgeBaseSet): The knowledge base.
            prompt (str): The prompt to be matched.
            top_k (int) = 10: The number of top matches to return.
            threshold_score (float) = 0.75: The threshold score for the top matches.

        Returns:
            list[tuple[float, str]]: The top matches with their scores and the corresponding chunks.
        """

        # Update the progress
        self.progress._update_progress(pss.KL_BASE_MATCHING, f"Matching knowledge base with prompt '{prompt}'.", {
            "knowledge_base": knowledge_base,
            "prompt": prompt,
            "top_k": top_k,
            "threshold_score": threshold_score
        })

        show_debug(f"Matching knowledge base with prompt '{prompt}'...")

        # Match the prompt with the knowledge base
        matched_results: list[tuple[float, str]] = knowledge_base.match_knowledge(self.embedding_model, prompt, top_k, threshold_score)

        # Update the progress
        self.progress._update_progress(pss.KL_BASE_MATCHED, f"Knowledge base matched with prompt '{prompt}' with {len(matched_results)} results.", {
            "knowledge_base": knowledge_base,
            "prompt": prompt,
            "top_k": top_k,
            "threshold_score": threshold_score,
            "matched_results": matched_results,
            "total_matched_results": len(matched_results)
        })

        show_debug(f"Knowledge base matched with prompt '{prompt}' with {len(matched_results)} results.")

        self.progress._update_progress(pss.IDLE)

        return matched_results