from typing import Dict, List, Optional

from kiss_ai_stack.core.ai_clients.ai_client_abc import AIClientAbc
from kiss_ai_stack.core.dbs.db_abc import VectorDBAbc
from kiss_ai_stack.core.models.core.rag_response import ToolResponse
from kiss_ai_stack.core.models.enums.tool_kind import ToolKind
from kiss_ai_stack.core.utilities.logger import LOG

from src.kiss_ai_stack.core.models.config.tool import ToolProperties


class Tool:
    """
    A unified interface for processing AI tasks with optional vector database support.

    This class handles query processing using an AI client and can optionally store and retrieve
    documents from a vector database. It ensures sensitive data, such as query answers, is protected
    and does not log sensitive content.
    """

    def __init__(
            self,
            properties: ToolProperties,
            ai_client: AIClientAbc,
            vector_db: Optional[VectorDBAbc] = None
    ):
        """
        Initialize a Tool instance.

        :param properties: Tool configurations object loaded from Yaml file.
        :param ai_client: The AI client instance used for query processing.
        :param vector_db: The optional vector database instance for storing documents. Defaults to None.
        """
        self.__properties = properties
        self.__ai_client = ai_client
        self.__vector_db = vector_db
        LOG.info(f'Tool initialized with kind: {self.tool_kind()}')

    def tool_kind(self) -> ToolKind:
        """
        Get the type of the tool.

        :returns: The type of the tool, such as RAG or PROMPT.
        """
        return self.__properties.kind

    def depth(self) -> int:
        """
        Returns the dept of retrieved documents to be considered.(RAG only)
        """
        return self.__properties.depth

    def dpr(self):
        """
        Returns whether document retrival uses dense retrival method. In any error occurred while dpr operation
        it will use the standard retrieval as fallback mode.
        """
        return self.__properties.dpr

    async def store_docs(self, documents: List[str], metadata_list: List[Dict]) -> List[str]:
        """
        Store documents and their metadata in the vector database.

        :param documents: A list of documents to be stored in the database.
        :param metadata_list: A list of metadata dictionaries corresponding to each document.
        :returns: A list of document IDs generated by the vector database.
        :raises IOError: If the vector database is not initialized or an error occurs while storing documents.
        """
        if self.__vector_db:
            LOG.info(f'Pipeline tool :: Storing {len(documents)} documents in the vector database.')
            try:
                ids = await self.__vector_db.push(documents=documents, metadata_list=metadata_list)
                LOG.info(f'Pipeline tool :: {len(ids)} documents successfully stored with generated IDs.')
                return ids
            except Exception as e:
                LOG.error(f'Pipeline tool :: Failed to store documents: {str(e)}')
                raise IOError(f'Pipeline tool :: Error storing documents: {str(e)}')
        else:
            error_message = 'Pipeline tool :: Vector database has not been initialized.'
            LOG.error(error_message)
            raise IOError(error_message)

    async def process_query(self, query: str) -> Optional[ToolResponse]:
        """
        Process a query using the AI client and optionally the vector database.

        If the tool is in retrieval-augmented generation (RAG) mode, it will retrieve relevant
        documents from the vector database before passing them to the AI client to generate a response.

        :param query: The query string to be processed by the tool.
        :returns: A response containing the generated answer, documents, metadata, and distances,
                  or None if no valid response is generated.
        :raises Exception: If an error occurs while processing the query, such as failure in retrieval or generation.
        """
        LOG.info(f'Pipeline tool :: Processing query with tool kind: {self.tool_kind()}')
        try:
            retrieved_docs = None

            if self.tool_kind() == ToolKind.RAG:
                LOG.info('Pipeline tool :: Performing retrieval-augmented generation (RAG) for query processing.')
                try:
                    if self.dpr():
                        query_embedding = await self.__ai_client.embed_text(query)
                        retrieved_docs = await self.__vector_db.dense_retrieve(query_embedding)
                except Exception as err:
                    LOG.error(f'Pipeline tool :: {err}')
                    LOG.warning(f'Pipeline tool :: DPR failed. Fallback to standard documents retrieval function')
                    retrieved_docs = await self.__vector_db.retrieve(query)

            if retrieved_docs:
                top_docs = retrieved_docs['documents'][:self.depth()]
                LOG.debug('Pipeline tool :: Chunk metadata retrieved but content not logged for security reasons.')
                answer = await self.__ai_client.generate_answer(query, top_docs)
                LOG.info('Pipeline tool :: Answer generated by AI client.')

                return ToolResponse(
                    answer=answer,
                    docs=top_docs,
                    metadata=retrieved_docs['metadatas']
                )
            else:
                LOG.info('Pipeline tool :: Using direct prompt mode for query processing.')
                answer = await self.__ai_client.generate_answer(query=query)
                LOG.info('Pipeline tool :: Answer generated by AI client (no documents).')

                return ToolResponse(answer=answer)

        except Exception as e:
            LOG.error(f'Pipeline tool :: Failed to process query: {str(e)}')
            raise e

    async def destroy(self, cleanup: bool = False):
        """
        Close connections and optionally clean up resources.

        :param cleanup: If True, cleans up the vector database collection (if initialized).
        """
        if cleanup and self.__vector_db is not None:
            LOG.info('Pipeline tool :: Cleaning up vector database collection.')
            await self.__vector_db.destroy()

        LOG.info('Pipeline tool :: Destroying AI client connection.')
        await self.__ai_client.destroy()
        return True
