import os
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from tools import ExtractTextFromImage, DescribeImage, read_excel, read_python, wiki_search, web_search, arxiv_search
from langchain_openai import ChatOpenAI
from typing import TypedDict, Annotated, Optional
from langfuse.callback import CallbackHandler


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]


class BasicAgent():
    def __init__(self, chat, vision_llm):
        extract_text_from_image = ExtractTextFromImage(vision_llm)
        describe_image = DescribeImage(vision_llm)

        self.tools = [extract_text_from_image.__call__, describe_image.__call__, read_excel, read_python, wiki_search, web_search, arxiv_search]
        self.chat_with_tools = chat.bind_tools(self.tools)
        self._initialize_graph()
        self._initialize_telemetry()
        print("BasicAgent initialized.")

    
    def _initialize_graph(self):
        builder = StateGraph(AgentState)

        # Define nodes
        builder.add_node("assistant", self.assistant)
        builder.add_node("tools", ToolNode(self.tools))

        # Define edges
        builder.add_edge(START, "assistant")
        builder.add_conditional_edges("assistant",tools_condition)
        builder.add_edge("tools", "assistant")

        # Compile the graph
        self.agent = builder.compile()

    def _initialize_telemetry(self):
        LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY")
        LANGFUSE_SECRET_KEY = os.getenv("LANGFUSE_SECRET_KEY")
        LANGFUSE_HOST = "https://cloud.langfuse.com"

        self.langfuse_handler = CallbackHandler(
            public_key=LANGFUSE_PUBLIC_KEY,
            secret_key=LANGFUSE_SECRET_KEY,
            host=LANGFUSE_HOST
        )

    
    def __call__(self, question: str, file_name : str) -> str:
        print(f"Agent received question: {question}. /nProvided file: {file_name}.")
        if file_name is not None and file_name!='':
            messages=[HumanMessage(content=f"{question}. The filename you have access to is {file_name}.")]
        else:
            messages=[HumanMessage(content=question)]
        response = self.agent.invoke({"messages":messages}, config={"callbacks": [self.langfuse_handler]})
        answer = response['messages'][-1].content
        print(f"Agent returning answer: {answer}")
        return answer

    
    def assistant(self, state: AgentState):
 
        sys_msg = SystemMessage(content=f"""
        You are a general AI assistant. I will ask you a question. Reason step by step and search for the information you need using available tools, one step at a time.
        If you do not have enough information to answer, use the tools available to search for it. 
        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.
        When providing the final answer, ONLY give [YOUR FINAL ANSWER]. Do not add anything else, no additional motivation or explanation, and do not return 'FINAL ANSWER:'.
        """)
        response = self.chat_with_tools.invoke([sys_msg] + state["messages"])
 
        return {
            "messages":  state["messages"] + [response],            
        }