import asyncio, json
from typing import List, Dict
from pathlib import Path

from animuz_core.pipelines.base import BasePipeline
from animuz_core.genai.base import BaseAgent
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.genai.tools import AnthropicRetrieverTool, OpenAIRetrieverTool

class AgenticRAG(BasePipeline):
    def __init__(self, embedding_client: EmbeddingClient, db_client: MyQdrantClient, agent: BaseAgent, tools: List[str] = [], unstructured_client: MyUnstructuredClient = None, azure_doc_ai_client: AzureDocAiClient = None) -> None:
        """
        Initializes the AgenticRAG object with the specified clients and models.

        Args:
            embedding_client (EmbeddingClient): The client for embedding data.
            db_client (MyQdrantClient): The client for the database.
            agent (BaseAgent): The LLM agent for the AgenticRAG object.
            tools (List[str]): The list of tool names available for the LLM agent.
            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 AgenticRAG object without returning anything.
        """
        super().__init__(embedding_client, db_client, unstructured_client, azure_doc_ai_client)

        self.agent = agent
        if tools:
            self.tools = self.create_tools(tools)
            self.agent.add_tools(self.tools)

    def create_tools(self, tool_names: List[str]):
        tools = dict()
        for tool_name in tool_names:
            if tool_name == "anthropic_retriever":
                tools["retriever"] = (AnthropicRetrieverTool.create_tool(), self.query_db)
            elif tool_name == "openai_retriever":
                tools["retriever"] = (OpenAIRetrieverTool.create_tool(), self.query_db)
        return tools

    async def add_doc(self, path: str, user_chat_id: str, is_chunked: bool) -> 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.
        """
        print("START parsing for doc", path)
        res = await self.parse_doc_markdown(path, is_chunked)
        print("done parsing for doc", path)
        properties = [{
                "text": data_row["text"],
                "source": data_row["source"],
                "fileDirectory": str(Path(path).parent),
                "filename": str(Path(data_row["metadata"]["filename"]).stem),
                "filetype": Path(data_row["metadata"]["filename"]).suffix,
                "userChatID": user_chat_id
            } for data_row in res]
        
        if len(res) < 1000:
            semaphore = asyncio.Semaphore(10)
            tasks = []
            emb_res = []
            for data_row in res:
                async with semaphore:
                    tasks.append(self.embed_single(str(data_row["text"])))
                    await asyncio.sleep(1)

            for task in tasks:
                emb_res.append(await task)

        else:
            emb_res = await self.batch_embed([str(data_row["text"]) for data_row in res])

        print("done embedding for doc", path)
        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)
        print("done uploading for doc", path)

    async def add_doc_from_s3(self, paths: List[str], user_chat_id: str, is_chunked: bool) -> None:
        semaphore = asyncio.Semaphore(10)
        tasks = []
        res = []
        for path in paths:
            async with semaphore:
                tasks.append(self.parse_doc_markdown(path, is_chunked))
                await asyncio.sleep(1)

        for task in tasks:
            res.extend(await task)

        properties = [{
            "text": data_row["text"],
            "source": data_row["source"],
            "fileDirectory": str(Path(path).parent),
            "filename": str(Path(data_row["metadata"]["filename"]).stem),
            "filetype": Path(data_row["metadata"]["filename"]).suffix,
            "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 delete_doc(self, filename: str, user_chat_id: str) -> None:
        conditions = [("fileName", str(Path(filename).stem)), 
                      ("fileType", Path(filename).suffix),
                      ("userChatID", user_chat_id)]
        await self.db_client.delete_doc(conditions= conditions)

    async def chat(self, messages: List[Dict[str, str]], user_chat_id: str) -> str:
        return await self.agent.get_reply_frontend(messages, user_chat_id)
    
    async def chat_two(self, messages: List[Dict[str, str]], user_chat_id: str, background_tasks) -> str:
        return await self.agent.get_reply_frontend_tool_use_ack(messages, user_chat_id, background_tasks)
    
    async def retrieve_result(self, chat_id: str) -> str:
        return await self.agent.get_chat_result(chat_id)
    
    async def streaming_chat(self, messages: List[Dict[str, str]], user_chat_id: str):
        async for chunk in self.agent.get_reply_frontend_streaming(messages, user_chat_id):
            yield json.dumps(chunk, ensure_ascii=False) + "\n"

    def clear_history(self) -> None:
        self.agent.clear_history()