from datetime import datetime, timezone
from colorama import Fore, Style  # type: ignore[import]
from langchain_core.tools import tool
from langchain.agents import create_agent
from agent.tools.search import web_search


@tool
def websearch_agent(query: str) -> str:
    """
    A web search agent that searches the internet and returns an answer.
    Use this tool when you need to find real-time or factual information
    from the web, such as current events, specific facts, or any
    knowledge that may require up-to-date sources.

    Args:
        query: The question or search query to look up on the web.
    """
    print(f"{Fore.YELLOW}[SupervisorAgent -> WebSearchAgent] {query}{Style.RESET_ALL}")
    base_agent = create_agent(
        model="google_genai:gemini-3-flash-preview",
        tools=[web_search],
        system_prompt=(
            f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
            f"Your memory are out of date. "
            f"All of truth that you believe without search are wrong. "
            f"You must search the web and find the lastest answer."
        ),
    )
    result = base_agent.invoke({"messages": [{"role": "user", "content": query}]})
    content = result["messages"][-1].content
    if isinstance(content, list):
        content = content[0].get("text", "")
    else:
        content = str(content)
    print(
        f"{Fore.YELLOW}[WebSearchAgent -> SupervisorAgent] {content}{Style.RESET_ALL}"
    )
    return content


if __name__ == "__main__":
    from dotenv import load_dotenv

    load_dotenv()
    answer = websearch_agent.invoke({"query": "What is LangGraph?"})
    print(answer)