from typing import TypedDict, Annotated, Optional
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage, ToolMessage
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.graph import START, StateGraph, END 
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
import os
from dotenv import load_dotenv
from tools import download_file_from_url, basic_web_search, extract_url_content, wikipedia_reader, transcribe_audio_file, question_youtube_video

# Load environment variables from .env file
load_dotenv()

OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
MAIN_LLM_MODEL = os.getenv("MAIN_LLM_MODEL", "google/gemini-2.0-flash-lite-001")

# Generate the chat interface, including the tools
if not OPENROUTER_API_KEY:
    raise ValueError("OPENROUTER_API_KEY is not set. Please ensure it is defined in your .env file or environment variables.")


def create_agent_graph():
    
    main_llm = ChatOpenAI(
        model=MAIN_LLM_MODEL, # e.g., "mistralai/mistral-7b-instruct"
        api_key=SecretStr(OPENROUTER_API_KEY), # Your OpenRouter API key
        base_url="https://openrouter.ai/api/v1", # Standard OpenRouter API base
        verbose=True # Optional: for debugging
    )


    tools = [download_file_from_url, basic_web_search, extract_url_content, wikipedia_reader, transcribe_audio_file, question_youtube_video] # Ensure these tools are defined
    chat_with_tools = main_llm.bind_tools(tools)

    class AgentState(TypedDict):
        messages: Annotated[list[AnyMessage], add_messages]
        file_url: Optional[str | None]
        file_ext: Optional[str | None]
        local_file_path: Optional[str | None]
        final_answer: Optional[str | None]
    
    def assistant(state: AgentState):
        return {
            "messages": [chat_with_tools.invoke(state["messages"])],
            "file_url": state.get("file_url", None),
            "file_ext": state.get("file_ext", None),
            "local_file_path": state.get("local_file_path", None),
            "final_answer": state.get("final_answer", None)
        }
    
    def file_path_updater_node(state: AgentState):
        download_tool_response = state["messages"][-1].content
        file_path = download_tool_response.split("Local File Path: ")[-1].strip()
        return {
            "local_file_path": file_path
        }
    
    def file_path_condition(state: AgentState) -> str:
        if state["messages"] and isinstance(state["messages"][-1], ToolMessage):
            tool_response = state["messages"][-1]
            if tool_response.name == "download_file_from_url":
                return "update_file_path"  # Route to file path updater if a file was downloaded
        return "assistant"  # Otherwise, continue with the assistant node

    def format_final_answer_node(state: AgentState) -> AgentState:
        """
        Formats the final answer based on the state.
        This node is reached when the assistant has completed its task.
        """
        final_answer = state["messages"][-1].content if state["messages"] else None
        if final_answer:
            state["final_answer"] = final_answer.split("FINAL ANSWER:")[-1].strip() #if FINAL_ANSWER isn't present we grab the whole string
        return state           
        
        
    # The graph
    builder = StateGraph(AgentState)
    
    builder.add_node("assistant", assistant)
    builder.add_edge(START, "assistant")
    builder.add_node("tools", ToolNode(tools))
    builder.add_node("file_path_updater_node", file_path_updater_node)
    builder.add_node("format_final_answer_node", format_final_answer_node)
    
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
        {
            "tools": "tools", 
            "__end__": "format_final_answer_node"  # This is the end node for the assistant
        }
    )
    builder.add_conditional_edges(
        "tools",
        file_path_condition,
        {
            "update_file_path": "file_path_updater_node",
            "assistant": "assistant"
        }
    )
    
    builder.add_edge("file_path_updater_node", "assistant")
    builder.add_edge("format_final_answer_node", END)
    graph = builder.compile()
    return graph

class BasicAgent:
    """
    A basic agent that can answer questions and download files.
    Requires a system message be defined in 'system_prompt.txt'.
    """
    def __init__(self, graph=None):

        with open("system_prompt.txt", "r", encoding="utf-8") as f:
            self.system_message = SystemMessage(content=f.read())
        
        if graph is None:
            self.graph = create_agent_graph()
        else:
            self.graph = graph

    def __call__(self, question: str, file_url: Optional[str] = None, file_ext: Optional[str] = None) -> str:
        """
        Call the agent with a question and optional file URL and extension.
        
        Args:
            question (str): The user's question.
            file_url (Optional[str]): The URL of the file to download.
            file_ext (Optional[str]): The file extension for the downloaded file.
        
        Returns:
            str: The agent's response.
        """
        if file_url and file_ext:
            question += f"\nREFERENCE FILE MUST BE RETRIEVED\nFile URL: {file_url}, File Extension: {file_ext}\nUSE A TOOL TO DOWNLOAD THIS FILE."
        state = {
            "messages": [self.system_message, HumanMessage(content=question)],
            "file_url": file_url,
            "file_ext": file_ext,
            "local_file_path": None,
            "final_answer": None
        }
        response = self.graph.invoke(state)
        for m in response["messages"]:
            m.pretty_print()
        return response["final_answer"] if response["final_answer"] else "No final answer generated."