import os
from pathlib import Path
from typing import List, Tuple

from animuz_core.pipelines.base import BasePipeline
from animuz_core.genai.base import BaseLLM
from animuz_core.embedding.embedding_client import EmbeddingClient
from animuz_core.ingest.custom_unstructured_client import MyUnstructuredClient
from animuz_core.ingest.azure_doc_ai_client import AzureDocAiClient
from animuz_core.vectordb.utils import MyQdrantClient

from animuz_core.utils import format_sources

class SimpleRAG(BasePipeline):
    def __init__(self, embedding_client: EmbeddingClient, db_client: MyQdrantClient, LLM: BaseLLM, unstructured_client: MyUnstructuredClient = None, azure_doc_ai_client: AzureDocAiClient = None) -> None:
        """
        Initializes the SimpleRAG object with the specified clients and models.

        Args:
            embedding_client (EmbeddingClient): The client for embedding data.
            db_client (MyQdrantClient): The client for the database.
            LLM (BaseLLM): The LLM model for the SimpleRAG object.
            unstructured_client (MyUnstructuredClient, optional): The client for unstructured data. If None, uses the default MyUnstructuredClient.
            azure_doc_ai_client (AzureDocAiClient, optional): The client for Azure Doc AI. If None, uses the default AzureDocAiClient.

        Returns:
            None: This function initializes the SimpleRAG object without returning anything.
        """
        super().__init__(embedding_client, db_client, unstructured_client, azure_doc_ai_client)

        self.llm = LLM

    async def add_doc(self, path: str, user_chat_id: str = "default") -> None:
        """
        Add a document to the database.

        Args:
            path (str): The path to the document.

        Returns:
            None: This function does not return anything.

        Raises:
            None: This function does not raise any exceptions.

        This function adds a document to the database by performing the following steps:
        1. Partition the document using the unstructured client or azure doc ai.
        2. Extract the properties of each partitioned data row.
        3. Get the embeddings for each partitioned data row using the embedding client.
        4. Extract the indices and values from the sparse embeddings.
        5. Upload the dense and sparse embeddings, indices, values, and properties to the database using the db client.

        Note: The document is expected to be in the specified path.
        """
        
        res = await self.parse_doc(path)

        properties = [{
                "text": data_row["text"],
                "source": data_row["source"],
                "file_directory": str(Path(path).parent),
                "filename": data_row["metadata"]["filename"],
                "filetype": data_row["metadata"]["filetype"],
                "userChatID": user_chat_id,
            } for data_row in res]
        
        emb_res = await self.batch_embed([str(data_row["text"]) for data_row in res])

        indices = [e["sparse_embedding"].keys() for e in emb_res]
        values = [e["sparse_embedding"].values() for e in emb_res]
            
        await self.upload_to_db([e["embedding"] for e in emb_res], indices, values, properties)

    async def query(self, query: str, user_chat_id: str) -> str:
        """
        Performs a query operation based on the input query string. Note that this is a stateless operation. No history is maintained.
        
        Args:
            query (str): The query string to search for.
            user_chat_id (str): The user chat ID used for multi-tenant retrieval.

        Returns:
            str: The response to the query.
        """
        search_results, _ = await self.query_db(query, user_chat_id)
        prompt = f"\n Use the following context to answer the user query: {format_sources(search_results)}. Be sure to include citations in your response (in the form of \"Source x: [answer]\" where x is an index). \n\n User query: {query}"
        reply = await self.llm.get_reply(prompt)
        self.llm.clear_history() #stateless chat
        return reply
