from typing import Optional
from collections.abc import Callable


class RagAgentModeEnum:
    VectorStore = "vectorstore"
    KG = "kg"
    API_BLAST = "api_blast"
    API_ONCOKB = "api_oncokb"


class RagAgent:
    def __init__(
        self,
        mode: str,
        model_name: Optional[str] = "gpt-3.5-turbo",
        connection_args: Optional[dict] = None,
        n_results: Optional[int] = 3,
        use_prompt: Optional[bool] = False,
        schema_config_or_info_dict: Optional[dict] = None,
        conversation_factory: Optional[Callable] = None,
        embedding_func: Optional[object] = None,
        documentids_workspace: Optional[list[str]] = None,
        agent_desc: Optional[str] = None,
    ) -> None:
        ######
        ##TO DO
        ######
        # mode: 'api' for the case where the agent is querying an API
        # use_prompt: default TRUE for self.mode == api

        """
        Create a RAG agent that can return results from a database or vector
        store using a query engine.

        Args:
            mode (str): The mode of the agent. Either "kg" or "vectorstore".

            model_name (str): The name of the model to use.

            connection_args (dict): A dictionary of arguments to connect to the
                database or vector store. Contains database name (in case of
                multiple DBs in one DBMS), host, port, user, and password.

            n_results: the number of results to return for method
                generate_response

            use_prompt (bool): Whether to use the prompt for the query. If
                False, will not retrieve any results and return an empty list.

            schema_config_or_info_dict (dict): A dictionary of schema
                information for the database. Required if mode is "kg".

            conversation_factory (Callable): A function used to create a
                conversation for creating the KG query. Required if mode is
                "kg".

            embedding_func (object): An embedding function. Required if mode is
                "vectorstore".

            documentids_workspace (Optional[List[str]], optional): a list of
                document IDs that defines the scope within which similarity
                search occurs. Defaults to None, which means the operations will
                be performed across all documents in the database.

        """
        self.mode = mode
        self.model_name = model_name
        self.use_prompt = use_prompt
        self.n_results = n_results
        self.documentids_workspace = documentids_workspace
        self.last_response = []
        self._agent_desc = agent_desc
        if self.mode == RagAgentModeEnum.KG:
            from .database_agent import DatabaseAgent

            if not connection_args:
                raise ValueError(
                    "Please provide connection args to connect to database."
                )

            if not schema_config_or_info_dict:
                raise ValueError("Please provide a schema config or info dict.")
            self.schema_config_or_info_dict = schema_config_or_info_dict

            self.agent = DatabaseAgent(
                model_name=model_name,
                connection_args=connection_args,
                schema_config_or_info_dict=self.schema_config_or_info_dict,
                conversation_factory=conversation_factory,
            )

            self.agent.connect()

            self.query_func = self.agent.get_query_results

        elif self.mode == RagAgentModeEnum.VectorStore:
            from .vectorstore_agent import VectorDatabaseAgentMilvus

            if not connection_args:
                raise ValueError(
                    "Please provide connection args to connect to vector store."
                )

            if not embedding_func:
                raise ValueError("Please provide an embedding function.")

            self.agent = VectorDatabaseAgentMilvus(
                embedding_func=embedding_func,
                connection_args=connection_args,
            )

            self.agent.connect()

            self.query_func = self.agent.similarity_search

        elif self.mode == RagAgentModeEnum.API_BLAST:
            from .api_agent.blast import (
                BlastFetcher,
                BlastInterpreter,
                BlastQueryBuilder,
            )
            from .api_agent.api_agent import APIAgent

            self.agent = APIAgent(
                conversation_factory=conversation_factory,
                query_builder=BlastQueryBuilder(),
                fetcher=BlastFetcher(),
                interpreter=BlastInterpreter(),
            )
            self.query_func = self.agent.execute
        elif self.mode == RagAgentModeEnum.API_ONCOKB:
            from .api_agent.oncokb import (
                OncoKBFetcher,
                OncoKBInterpreter,
                OncoKBQueryBuilder,
            )
            from .api_agent.api_agent import APIAgent

            self.agent = APIAgent(
                conversation_factory=conversation_factory,
                query_builder=OncoKBQueryBuilder(),
                fetcher=OncoKBFetcher(),
                interpreter=OncoKBInterpreter(),
            )
            self.query_func = self.agent.execute
        else:
            raise ValueError(
                "Invalid mode. Choose either 'kg', 'vectorstore', 'api_blast', or 'api_oncokb'."
            )

    @property
    def agent_description(self):
        return self._agent_desc

    @agent_description.setter
    def agent_description(self, val: Optional[str] = None):
        self._agent_desc = val

    def generate_responses(self, user_question: str) -> list[tuple]:
        """
        Run the query function according to the mode and return the results in a
        uniform format (list of tuples, where the first element is the text for
        RAG and the second element is the metadata).

        Args:
            user_question (str): The user question.

        Returns:
            results (List[tuple]): A list of tuples containing the results.

        Todo:
            Which metadata are returned?
        """
        self.last_response = []
        if not self.use_prompt:
            return []
        if self.mode == RagAgentModeEnum.KG:
            results = self.query_func(user_question, self.n_results)
            response = [
                (
                    result.page_content,
                    result.metadata,
                )
                for result in results
            ]
        elif self.mode == RagAgentModeEnum.VectorStore:
            results = self.query_func(
                user_question,
                self.n_results,
                doc_ids=self.documentids_workspace,
            )
            response = [
                (
                    result.page_content,
                    result.metadata,
                )
                for result in results
            ]
        elif self.mode in [
            RagAgentModeEnum.API_BLAST,
            RagAgentModeEnum.API_ONCOKB,
        ]:
            final_answer = self.query_func(user_question)
            if final_answer is not None:
                response = [("response", final_answer)]
            else:
                response = [("error", final_answer)]

        else:
            raise ValueError(
                "Invalid mode. Choose either 'kg', 'vectorstore', 'api_blast', "
                "or 'api_oncokb'."
            )
        self.last_response = response
        return response

    def get_description(self):
        if self.agent_description is not None:
            return self.agent_description
        if self.mode == RagAgentModeEnum.KG:
            return self.agent.get_description()
        elif self.mode == RagAgentModeEnum.VectorStore:
            return self.agent.get_description(self.documentids_workspace)
        elif self.mode == RagAgentModeEnum.API_BLAST:
            tool_name = "BLAST"
            tool_desc = (
                "The Basic Local Alignment Search Tool (BLAST) "
                "finds regions of local similarity between sequences. "
                "BLAST compares nucleotide or protein sequences to "
                "sequence databases and calculates the statistical "
                "significance of matches."
            )
            return self.agent.get_description(tool_name, tool_desc)
        elif self.mode == RagAgentModeEnum.API_ONCOKB:
            tool_name = "OncoKB"
            tool_desc = (
                "OncoKB is a precision oncology knowledge base "
                "and contains information about the effects "
                "and treatment implications of specific cancer gene alterations."
            )
            return self.agent.get_description(tool_name, tool_desc)
        else:
            raise ValueError(
                "Invalid mode. Choose either 'kg', 'vectorstore', 'api_blast', "
                "or 'api_oncokb'."
            )
