import os
from datetime import datetime, timezone
from dotenv import load_dotenv
from colorama import Fore, Style
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage
from agent.tools.math_solver import math_solver
from agent.agents.websearch import websearch_agent

load_dotenv()


def supervisor_agent():
    """Return a supervisor agent instance with math_solver and websearch_agent."""
    return create_agent(
        model="google_genai:gemini-3-flash-preview",
        tools=[math_solver, websearch_agent],
        system_prompt=(
            f"You are a supervisor agent. "
            f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
            f"Your memory are out of date. "
            f"For math or calculation questions, use the math_solver tool. "
            f"For questions that need real-time, use the websearch_agent tool. "
            f"Provide a concise and accurate final answer."
        ),
    )


def run(query: str) -> str:
    """Entry point: let the supervisor agent finish the work."""
    print(f"{Fore.CYAN}[Supervisor] Processing query...{Style.RESET_ALL}")
    agent = supervisor_agent()
    result = agent.invoke({"messages": [HumanMessage(content=query)]})
    content = result["messages"][-1].content
    if isinstance(content, list):
        return content[0].get("text", "")
    return str(content)


if __name__ == "__main__":
    agent = supervisor_agent()
    chat_history: list = []
    while True:
        query = input("\nYou: ")
        if query.lower() in ("exit", "quit"):
            break
        chat_history.append(HumanMessage(content=query))
        result = agent.invoke({"messages": chat_history})
        chat_history = result["messages"]
        content = chat_history[-1].content
        if isinstance(content, list):
            content = content[0].get("text", "")
        print(f"Agent: {content}")