# flake8: noqa
import requests
import pandas as pd

from io import StringIO
from decouple import config
from tempfile import NamedTemporaryFile
from langchain_community.tools import BaseTool
from llama import Context, LLMEngine, Type
from app.vectorstores.base import VectorStoreMain
from app.datasource.loader import DataLoader
from prisma.models import Datasource

from langchain.agents.agent_types import AgentType
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
from langchain_openai import ChatOpenAI


class DatasourceFinetuneTool(BaseTool):
    name = "datasource"
    description = "useful for when you need to answer questions"
    return_direct = False

    def _run(
        self,
        question: str,
    ) -> str:
        """Use the tool."""

        class Question(Type):
            question: str = Context("A question")

        class Answer(Type):
            answer: str = Context("An answer to the question")

        llm = LLMEngine(
            id=self.metadata["agent_id"],
            config={"production.key": config("LAMINI_API_KEY")},
            model_name="chat/gpt-3.5-turbo",
        )
        input = Question(question=question)
        output = llm(input=input, output_type=Answer)
        return output.answer

    async def _arun(
        self,
        question: str,
    ) -> str:
        """Use the tool asynchronously."""

        class Question(Type):
            question: str = Context("A question")

        class Answer(Type):
            answer: str = Context("An answer to the question")

        llm = LLMEngine(
            id=self.metadata["agent_id"],
            config={"production.key": config("LAMINI_API_KEY")},
            model_name="chat/gpt-3.5-turbo",
        )
        input = Question(question=question)
        output = llm(input=input, output_type=Answer)
        return output.answer


class DatasourceTool(BaseTool):
    name = "datasource"
    description = "useful for when you need to answer questions"
    return_direct = False

    def _run(
        self,
        question: str,
    ) -> str:
        """Use the tool."""
        vector_store = VectorStoreMain(
            options=self.metadata["options"],
            vector_db_provider=self.metadata["provider"],
            embeddings_model_provider=self.metadata["embeddings_model_provider"],
        )
        result = vector_store.query_documents(
            prompt=question,
            datasource_id=self.metadata["datasource_id"],
            query_type=self.metadata["query_type"],
            top_k=3,
        )
        return result

    async def _arun(
        self,
        question: str,
    ) -> str:
        """Use the tool asynchronously."""
        vector_store = VectorStoreMain(
            options=self.metadata["options"],
            vector_db_provider=self.metadata["provider"],
            embeddings_model_provider=self.metadata["embeddings_model_provider"],
        )
        result = vector_store.query_documents(
            prompt=question,
            datasource_id=self.metadata["datasource_id"],
            query_type=self.metadata["query_type"],
            top_k=3,
        )
        return result


class StructuredDatasourceTool(BaseTool):
    name = "structured datasource"
    description = "useful for when need answer questions"
    return_direct = False

    def _load_xlsx_data(self, datasource: Datasource):
        with NamedTemporaryFile(suffix=".xlsx", delete=True) as temp_file:
            if datasource.url:
                response = requests.get(datasource.url)
                temp_file.write(response.content)
            else:
                temp_file.write(datasource.content)
            temp_file.flush()
            df = pd.read_excel(temp_file.name, engine="openpyxl")
        return df

    def _load_csv_data(self, datasource: Datasource):
        if datasource.url:
            response = requests.get(datasource.url)
            file_content = StringIO(response.text)
        else:
            file_content = StringIO(datasource.content)
        df = pd.read_csv(file_content)
        return df

    def _run(
        self,
        question: str,
    ) -> str:
        """Use the tool."""
        datasource: Datasource = self.metadata["datasource"]
        if datasource.type == "CSV":
            df = self._load_csv_data(datasource)
        elif datasource.type == "XLSX":
            df = self._load_xlsx_data(datasource)
        else:
            data = DataLoader(datasource=datasource).load()
            df = pd.DataFrame(data)
        agent = create_pandas_dataframe_agent(
            ChatOpenAI(
                temperature=0,
                model="gpt-4-0613",
                openai_api_key=config("OPENAI_API_KEY"),
            ),
            df,
            verbose=True,
            agent_type=AgentType.OPENAI_FUNCTIONS,
        )
        output = agent.run(question)
        return output

    async def _arun(
        self,
        question: str,
    ) -> str:
        """Use the tool asynchronously."""
        datasource: Datasource = self.metadata["datasource"]
        if datasource.type == "CSV":
            df = self._load_csv_data(datasource)
        elif datasource.type == "XLSX":
            df = self._load_xlsx_data(datasource)
        else:
            data = DataLoader(datasource=datasource).load()
            df = pd.DataFrame(data)
        agent = create_pandas_dataframe_agent(
            ChatOpenAI(
                temperature=0,
                model="gpt-4-0613",
                openai_api_key=config("OPENAI_API_KEY"),
            ),
            df,
            verbose=True,
            agent_type=AgentType.OPENAI_FUNCTIONS,
        )
        output = await agent.arun(question)
        return output
