# agente_functions.py

from typing import TypedDict, List
from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain.agents import load_tools
from langgraph.prebuilt import create_react_agent

# --- Define a estrutura de estado do grafo ---
class AgentState(TypedDict):
    user_query: str
    router_decision: str
    web_answer: str
    scientific_answer: str
    final_answer: str

# --- Ferramentas e Funções dos Agentes ---

@tool
def busca_web(query: str) -> list:
  """
  Busca na web por um termo específico usando Tavily Search.
  """
  tavily_search = TavilySearchResults(max_results=2, search_depth='advanced', max_tokens=1000)
  resultado_busca = tavily_search.invoke(query)
  return resultado_busca

def funcao_agente_web(state: AgentState, llm, tools: List) -> dict:
  """
  Executa um agente ReAct que realiza uma busca na web.
  """
  system_prompt = """Atue como um assistente útil.
    Use as ferramentas fornecidas para responder às perguntas do usuário.
    - busca_web: Retorna os resultados de uma busca na web.
    Use a busca_web sempre que o usuário fizer uma pergunta sobre um tema específico e retorne o link dos artigos na resposta.
  """
  agente_web = create_react_agent(model=llm, tools=tools, prompt=system_prompt)
  resultado = agente_web.invoke({"messages": [("user", state["user_query"])]})
  resposta_final = resultado['messages'][-1].content
  return {"web_answer": resposta_final}

def funcao_agente_cientifico(state: AgentState, llm, tools: List) -> dict:
    """
    Executa um agente ReAct que realiza uma busca no Arxiv.
    """
    system_prompt = """Atue como um assistente útil.
    Use as ferramentas fornecidas para responder às perguntas do usuário.
    - tool_cientifica: Retorna os resultados de uma busca no arxiv.
    Use a tool_cientifica sempre que o usuário fizer uma pergunta sobre um tema específico e retorne o título dos artigos na resposta.
    """
    agente_cientifico = create_react_agent(llm, tools, prompt=system_prompt)
    resultado = agente_cientifico.invoke({"messages": [("user", state["user_query"])]})
    resposta_final = resultado['messages'][-1].content
    return {"scientific_answer": resposta_final}

def router_agent(state: AgentState, llm) -> dict:
    """
    Agente roteador que decide qual ferramenta usar com base na query do usuário.
    """
    router_prompt_template = """Você é um agente roteador. Sua tarefa é decidir qual o agente é mais apropriado para responder à pergunta do usuário.
    Se a pergunta mencionar explicitamente "fontes da web", "notícias" ou for de natureza geral, escolha 'web_search'.
    Se a pergunta mencionar "artigos científicos", "pesquisas", "estudos" ou for de natureza técnica/acadêmica, escolha 'scientific_search'.
    Para perguntas gerais sem especificação, prefira 'web_search'.

    Pergunta do usuário: {user_query}

    Responda APENAS com o nome da opção escolhida: 'web_search' ou 'scientific_search'.
    """
    prompt = router_prompt_template.format(user_query=state["user_query"])
    response = llm.invoke(prompt)
    decision = response.content.strip().lower()

    if "scientific_search" in decision:
        return {"router_decision": "scientific_search"}
    else:
        return {"router_decision": "web_search"}

def supervisor_node(state: AgentState) -> dict:
    """
    Combina os resultados dos outros nós e formata a resposta final.
    """
    web_results = state.get("web_answer", "Não foi realizada pesquisa na Web.")
    scientific_results = state.get("scientific_answer", "Não foi realizada pesquisa por artigos científicos.")

    # Constrói a resposta final baseada nos resultados disponíveis
    final_answer = "## Resultados da Pesquisa\n\n"
    if state.get("web_answer"):
        final_answer += "### Resultados da Pesquisa na Web\n"
        final_answer += web_results
    if state.get("scientific_answer"):
        final_answer += "\n\n### Resultados de Artigos Científicos\n"
        final_answer += scientific_results

    return {"final_answer": final_answer}