from tools import (
    general_tools, 
    file_agent_tools, 
    data_agent_tools, 
    math_agent_tools, 
    analyze_video_tools, 
    youtube_transcript_tools, 
    google_search,
    wiki_search,
    arxiv_search
)
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from langgraph_supervisor import create_supervisor


llm = ChatOpenAI(model="o4-mini")

memory = MemorySaver()

with open("system_prompt.txt", "r") as f:
    prompt = f.read()
    
general_agent = create_react_agent(
    model=llm,
    tools=general_tools(),
    checkpointer=memory,
    prompt=prompt
)

# Create agents
file_agent = create_react_agent(
    model=llm,
    tools=file_agent_tools(),
    name="file_reader",
    prompt="You read files. Use tools to read files."
)

math_agent = create_react_agent(
    model=llm,
    tools=math_agent_tools(),
    name="calculator", 
    prompt="You do math. Use tools for all calculations."
)

data_agent = create_react_agent(
    model=llm,
    tools=data_agent_tools(),
    name="data_processor",
    prompt="You process data. Use tools to filter and extract data."
)

# Create video analysis agents
video_agent = create_react_agent(
    model=llm,
    tools=analyze_video_tools(),
    name="video_analyzer",
    prompt="""You analyze visual content in videos. Use tools to detect and track objects. 
    The object_detection tool is a general object detection model. Use this for general cases.
    The analyze_video_content uses both the object detection model and a vision llm to analyze frames with content given a question. 
    Use this for more difficult questions."""
)

transcript_agent = create_react_agent(
    model=llm,
    tools=youtube_transcript_tools(),
    name="transcript_analyzer", 
    prompt="You analyze audio/speech content in videos. Use tools to get transcripts."
)

wiki_agent = create_react_agent(
    model=llm,
    tools=[wiki_search],
    name="wiki_analyst", 
    prompt="You search information from wikipedia."
)

google_agent = create_react_agent(
    model=llm,
    tools=[google_search],
    name="google_search_analyst", 
    prompt="You search information from google search."
)

arxiv_agent = create_react_agent(
    model=llm,
    tools=[arxiv_search],
    name="arxiv_analyst", 
    prompt="You search information from arxiv."
)

excel_prompt = """You are a supervisor. You coordinate file_reader, calculator, and data_processor to solve problems step by step. 
Do not do calculations or file reading yourself, use the tools.
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.
"""

video_analyzer_prompt = """You coordinate video_analyzer and transcript_analyzer to answer questions about YouTube videos. 
Use video_analyzer for visual questions (objects, people, actions). Use transcript_analyzer for audio questions (what people say).
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.
"""

search_analyzer_prompt = """You coordinate different search agents to answer questions. 
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.
"""

# Supervisor
excel_supervisor = create_supervisor(
    [file_agent, math_agent, data_agent],
    model=llm,
    prompt=excel_prompt
).compile()

# Video supervisor
video_supervisor = create_supervisor(
    [video_agent, transcript_agent],
    model=llm,
    prompt=video_analyzer_prompt
).compile()

# search supervisor
search_supervisor = create_supervisor(
    [wiki_agent, google_agent, arxiv_agent],
    model=llm,
    prompt=search_analyzer_prompt
).compile()