from langchain_text_splitters import CharacterTextSplitter
from utils import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from utils import ChatOpenAI
import pandas as pd
from tqdm import tqdm
from langchain_community.document_loaders import WebBaseLoader
from langchain.tools.retriever import create_retriever_tool
from langchain import hub
from langchain.agents import create_openai_tools_agent
from langchain.agents.agent import AgentExecutor

from config_parameters import ConfigParameters


class HW:
    def __init__(self):
        config_parameters = ConfigParameters()
        llm = ChatOpenAI(temperature=0.0, course_api_key=config_parameters.course_api_key)

        self.df = pd.read_csv('gannibal.csv')
        url = 'https://ru.wikipedia.org/wiki/Ганнибал,_Абрам_Петрович'

        loader = WebBaseLoader(url)
        documents = loader.load()

        # Определяем сплиттер
        splitter = CharacterTextSplitter(
            separator="\n\n",  # символ-разделитель, по умолчанию переход к новому абзацу '\n\n'
            chunk_size=2800,  # размер документа в символах
            chunk_overlap=100,  # насколько соседние документы могут перекрывать друг-друга
            length_function=len,  # функция, по которой считается размер документа
            is_separator_regex=False,  # является ли разделитель регулярным выражением
        )

        split_texts = splitter.split_documents(documents)

        # Задаём embedding model
        embeddings = OpenAIEmbeddings(course_api_key=config_parameters.course_api_key)

        # Создаём векторное хранилище
        db = FAISS.from_documents(split_texts[:4], embeddings)

        # Задаём ретривер
        retriever = db.as_retriever(
            search_type="similarity",  # тип поиска похожих документов
            k=4,  # количество возвращаемых документов (Default: 4)
            score_threshold=None,  # минимальный порог для поиска "similarity_score_threshold"
        )

        tool = create_retriever_tool(
            retriever,  # наш ретривер
            "search_web",  # имя инструмента
            "Searches and returns data from page",  # описание инструмента подается в ЛЛМ
        )

        tools = [tool]

        prompt = hub.pull("hwchase17/openai-tools-agent", api_key=config_parameters.langchain_api_key)
        prompt += """You must give a structured answer in Russian language in exactly the right format. Extra words not needed. Your answer must be less than 70 symbols."""

        agent = create_openai_tools_agent(llm, tools, prompt)
        self.agent_executor = AgentExecutor(agent=agent, tools=tools)

    def run(self):
        ans = []
        for task in tqdm(self.df['question']):
            answer = self.agent_executor.invoke({"input": task})["output"][:70]
            print(f'question: {task}; \nanswer: {answer}')
            ans.append(answer)

        self.df['answer'] = ans
        self.df.to_csv('solution.csv', index=False)


if __name__ == "__main__":
    hw = HW()
    hw.run()