from llama_index.core.tools import FunctionTool
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.tools.wikipedia import WikipediaToolSpec
from langfuse.llama_index import LlamaIndexInstrumentor
from llama_index.llms.ollama import Ollama
from llama_index.core.agent.workflow import FunctionAgent

from multimodality_tools import get_image_qa_tool, get_transcription_tool, \
    get_excel_analysis_tool, get_excel_tool, get_csv_analysis_tool, get_csv_tool

class BasicAgent:
    def __init__(self, ollama=False, langfuse=False):
        if not ollama:
            llm = HuggingFaceInferenceAPI(model_name="Qwen/Qwen2.5-Coder-32B-Instruct")
        else:
            llm = Ollama(model="mistral:latest", request_timeout=120.0)

        # Langfuse
        self.langfuse = langfuse
        if self.langfuse:
            self.instrumentor = LlamaIndexInstrumentor()
            self.instrumentor.start()

        # Initialize tools
        tool_spec = DuckDuckGoSearchToolSpec()
        search_tool = FunctionTool.from_defaults(tool_spec.duckduckgo_full_search)

        # Convert into a LoadAndSearchToolSpec because the wikipedia search tool returns
        # entire Wikipedia pages and this can pollute the context window of the LLM
        wiki_spec = WikipediaToolSpec()
        wiki_search_tool = wiki_spec.to_tool_list()[1]

        # Convert into a LoadAndSearchToolSpec because the wikipedia search tool returns
        # entire Wikipedia pages and this can pollute the context window of the LLM
        # TODO this does not work so well. We need to make the retriever return the top 5 chunks or sth.
        # wiki_search_tool_las = LoadAndSearchToolSpec.from_defaults(wiki_search_tool).to_tool_list()

        self.agent = FunctionAgent(
            tools=[search_tool, wiki_search_tool, get_image_qa_tool(),
                   get_transcription_tool(), get_excel_analysis_tool(), get_excel_tool(),
                   get_csv_analysis_tool(), get_csv_tool()],
            llm=llm,
            verbose=True,
            system_prompt = (
                "You are a general AI assistant. I will ask you a question. "
                "Report your thoughts, and finish your answer with the following template: "
                "FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number "
                "OR as few words as possible OR a comma separated list of numbers and/or "
                "strings. If you are asked for a number, don't use comma to write your "
                "number neither use units such as $ or percent sign unless specified otherwise. "
                "If you are asked for a string, don't use articles, neither abbreviations (e.g. "
                "for cities), and write the digits in plain text unless specified otherwise. If "
                "you are asked for a comma separated list, apply the above rules depending of "
                "whether the element to be put in the list is a number or a string."
            )
        )

        # self.ctx = Context(self.agent)

    async def __call__(self, question: str, task_id: str = None) -> str:
        file_str = ""
        if task_id:
            file_str = f'\nIf you need to load a file, do so by providing the id "{task_id}".'

        response = await self.agent.run(user_msg=question + file_str) # ctx=self.ctx)

        if self.langfuse:
            self.instrumentor.flush()

        return response.response.content.replace("FINAL ANSWER:", "").strip()