import logging
import os

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.chat_engine.types import AgentChatResponse
from llama_index.llms.openai import OpenAI
from llama_index.core import StorageContext, load_index_from_storage

logger = logging.getLogger("app.llama_index")


class Assistant:
    def __init__(
        self,
        data_dir: str = "data/pdf",
        model_name: str = "gpt-4",
    ) -> None:
        self.data_dir = data_dir
        self.model_name = model_name

        storage_dir = os.path.join(self.data_dir, "storage")
        if os.path.isdir(storage_dir):
            logger.info("Loading index from storage")
            storage_context = StorageContext.from_defaults(persist_dir=storage_dir)
            self.index = load_index_from_storage(storage_context)
        else:
            self.documents = SimpleDirectoryReader(self.data_dir).load_data()
            self.index = VectorStoreIndex(self.documents)
            self.index.storage_context.persist(storage_dir)

        llm = OpenAI(model=self.model_name)
        self.chat_engine = self.index.as_chat_engine(
            chat_mode="openai",
            llm=llm,
            verbose=True,
            system_prompt=(
                "Du er en assistent til at hjælpe med at besvare spørgsmål omkring årsopgørelsen fra skat."
                " Din viden er baseret på årsopgørelsen og relaterede emner fra skat.dk."
                " Du kan svare på spørgsmål omkring årsopgørelsen, og du kan også stille spørgsmål til brugeren for at få mere information."
                " Du kan også bede brugeren om at uddybe, hvis du ikke forstår spørgsmålet."
                " Du skal opgive de kilder, du bruger til at besvare spørgsmålet, og du skal også oplyse, hvis du ikke kan besvare spørgsmålet."
            ),
        )

    def __call__(self, question: str, history: list[str]) -> str:
        response: AgentChatResponse = self.chat_engine.chat(
            question,
            tool_choice="query_engine_tool",
        )
        return response.response