from langchain_tavily import TavilySearch
from langchain_community.document_loaders import WikipediaLoader
from langchain_ollama import ChatOllama
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor

llm_model=ChatOllama(model="qwen2:7b", temperature=0.0, max_tokens=1000, base_url="http://localhost:11434")

# Tool for searching wikipedia and returning the content
def search_wikipedia(query: str) -> str:
    """Searches Wikipedia for the given query and returns the first result's content.
    Args:
        query (str): The search query to use for Wikipedia. Should be a simple keyword or phrase. Do not enter the URL or any other complex query.
    """

    loader = WikipediaLoader(query=query, load_max_docs=1)
    docs = loader.load()
    if docs:
        return "From: " + docs[0].metadata['source'] + "\n\nContent: " + docs[0].page_content
    return "No content found."

# Tool for searching the web using Tavily
def search_web(query: str) -> str:
    """Searches the web for the given query and returns the first result's content.
    Args:
        query (str): The search query to use for web search. Should be a simple keyword or phrase. Do not enter the URL or any other complex query.
    """
    results = TavilySearch(max_results=1).invoke(input=query)
    if results:
        return "From: " + results['results'][0]['url'] + "\n\nContent: " + results['results'][0]['content']
    return "No content found."

research_agent = create_react_agent(
    model=llm_model,
    tools=[search_wikipedia, search_web],
    prompt=(
        "You are a research agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with research-related tasks, DO NOT do any math\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    ),
    name="research_agent",
)

def add(a: int, b: int) -> int:
    """
        Add two integers and return the result.
        Inputs:
            a: first integer
            b: second integer
    """
    return a + b

def subtract(a: int, b: int) -> int:
    """
        Subtract two integers and return the result.
        Inputs:
            a: first integer
            b: second integer
    """
    return a - b

def multiply(a: int, b: int) -> int:
    """
        Multiply two integers and return the result.
        Inputs:
            a: first integer
            b: second integer
    """
    return a * b

def divide(a: int, b: int) -> float:
    """
        Divide two integers and return the result.
        Inputs:
            a: first integer
            b: second integer
        Note: If b is 0, return "Error: Division by zero is not allowed."
    """
    if b == 0:
        return "Error: Division by zero is not allowed."
    return a / b

def modulo(a: int, b: int) -> int:
    """
        Calculate the modulo of two integers and return the result.
        Inputs:
            a: first integer
            b: second integer
        Note: If b is 0, return "Error: Division by zero is not allowed."
    """
    if b == 0:
        return "Error: Division by zero is not allowed."
    return a % b

math_agent = create_react_agent(
    model=llm_model,
    tools=[add, subtract, multiply, divide, modulo],
    prompt=(
        "You are a math agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with math-related tasks\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    ),
    name="math_agent",
)
def create_agent():
    supervisor = create_supervisor(
        model=llm_model,
        agents=[research_agent, math_agent],
        prompt=(
            "You are a supervisor managing two agents:\n"
            "- a research agent. Assign research-related tasks to this agent\n"
            "- a math agent. Assign math-related tasks to this agent\n"
            "Assign work to one agent at a time, do not call agents in parallel.\n"
            "Do not do any work yourself. When you are ready to give the final answer, do so as simply as possible.If the answer is a number, write only the number do not use commas or any special characters or any other text. \n"
            "If the answer is a string, write only the string without any additional text.\n"
            "If the answer is a list, write the items in the list separated by commas while following the above rules for numbers and string where appropriate.\n"
            "Preface your answer with FINAL ANSWER: {answer}\n"
            "Example: How many planets are in the solar system?\n Answer: 8\n"
            "Example: What is the capital of France?\n Answer: Paris\n"
            "Example: What is 2 + 2?\n Answer: 4\n"
        ),
        add_handoff_back_messages=True,
        output_mode="full_history",
    ).compile()

    return supervisor