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 langgraph.errors import GraphRecursionError
from agent.api.api import get_llm
from agent.tools.search import web_search


@tool
def websearch_agent(query: str) -> str:
    """
    A single 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.

    Pros:
    - Has continuous memory across search steps, allowing deep investigation on a single topic.
    Cons:
    - Narrow field of view, can only follow one search thread at a time.
    - May fail after too many steps due to token limit overflow.

    Prefer websearch_agents for complex questions requiring broad, multi-source research.
    Use this tool for simple, direct factual lookups.

    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=get_llm(),
        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."
            f"Just run 1 turn search. "
        ),
    )
    try:
        result = base_agent.invoke(
            {"messages": [{"role": "user", "content": query}]},
            # config={"recursion_limit": 10},
        )
        content = result["messages"][-1].content
        if isinstance(content, list):
            content = content[0].get("text", "")
        else:
            content = str(content)
    except GraphRecursionError:
        print(
            f"{Fore.RED}[WebSearchAgent] Recursion limit reached, returning partial results.{Style.RESET_ALL}"
        )
        content = "Search completed but no definitive answer was found within the allowed steps."
    except Exception as e:
        error_msg = str(e)
        print(f"{Fore.RED}[WebSearchAgent] Error: {error_msg}{Style.RESET_ALL}")
        content = (
            f"Search agent failed with error: {error_msg}. "
            f"Recommend retrying with the web_search_agents tool to avoid context length overflow."
        )
    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)