# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, MessagesState
import requests
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
import os
from dotenv import load_dotenv
from .tools import tools
from typing import Any

load_dotenv() 

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

class BasicAgent:
    def __init__(self):
         # Load environment variables from .env file
        self.model = ChatOpenAI(model_name="gpt-4o", temperature=0.0, api_key=os.getenv("OPENAI_API_KEY"))
        with open("src/prompts/system.txt", "r") as file:
            self.system_prompt = file.read()

        builder = StateGraph(MessagesState)
        builder.add_node("assistant", self.assistant)
        builder.add_node("tools", ToolNode(tools=tools))
        builder.add_edge(START, "assistant")
        builder.add_conditional_edges("assistant", tools_condition)
        builder.add_edge("tools", "assistant")
        self.graph = builder.compile()
        print("BasicAgent initialized.")
    
    def __call__(self, question: str) -> str:
        print(f"Agent received question...")
        message = HumanMessage(content=question)
        answer = self.graph.invoke({"messages": [message]})
        last_message = answer["messages"][-1]
        if isinstance(last_message, AIMessage) and "FINAL ANSWER:" in last_message.content:
            content = last_message.content
            # Find the position after "FINAL ANSWER: " and return only the text after it
            final_answer_pos = content.find("FINAL ANSWER:") + len("FINAL ANSWER:")
            return content[final_answer_pos:].strip()
            
        return "I'm sorry, I can't answer that question."
    
    def download_file(self, task_id: str) -> Any:
        api_url = DEFAULT_API_URL
        file_url = f"{api_url}/files/{task_id}"
        response = requests.get(file_url)
        return response.content
    
    def assistant(self, state: MessagesState):
        """
        This function takes a question and returns a list of sub-tasks.
        """
        sys_msg = SystemMessage(content=self.system_prompt)
        llm_with_tools = self.model.bind_tools(tools, parallel_tool_calls=False)
        state["messages"] = [sys_msg] + state["messages"]
        return { "messages": [llm_with_tools.invoke(state["messages"])] }