# main.py

import os
import gradio as gr
from functools import partial

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import load_tools
from langgraph.graph import StateGraph, START, END
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool

# Importa as funções e a classe de estado do outro arquivo
from agente_functions import (
    AgentState,
    busca_web,
    funcao_agente_web,
    funcao_agente_cientifico,
    supervisor_node,
    router_agent
)

# --- 1. Configuração do Ambiente ---
# É recomendado carregar as chaves de um arquivo .env para segurança
from dotenv import load_dotenv
load_dotenv()
api_key = os.getenv("GOOGLE_API_KEY")
tavily_api_key = os.getenv("TAVILY_API_KEY")

# Para este exemplo, vamos supor que as chaves estão como variáveis de ambiente
os.environ["GOOGLE_API_KEY"] = api_key
os.environ["TAVILY_API_KEY"] = tavily_api_key

if not os.environ.get("GOOGLE_API_KEY") or not os.environ.get("TAVILY_API_KEY"):
    print("Por favor, defina as variáveis de ambiente GOOGLE_API_KEY e TAVILY_API_KEY.")
    exit()

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

# --- 2. Inicialização do Modelo e Ferramentas ---
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
web_tools = [busca_web]
scientific_tools = load_tools(["arxiv"])


# --- 3. Construção do Grafo com LangGraph ---
workflow = StateGraph(AgentState)

# Usando functools.partial para passar argumentos fixos (llm, tools) para as funções dos nós
web_search_node = partial(funcao_agente_web, llm=llm, tools=web_tools)
scientific_search_node = partial(funcao_agente_cientifico, llm=llm, tools=scientific_tools)
router_node = partial(router_agent, llm=llm)

# Adiciona os nós ao grafo
workflow.add_node("router", router_node)
workflow.add_node("web_search", web_search_node)
workflow.add_node("scientific_search", scientific_search_node)
workflow.add_node("supervisor", supervisor_node)

# Define as arestas (fluxo de trabalho)
workflow.add_edge(START, "router")
workflow.add_conditional_edges(
    "router",
    lambda state: state['router_decision'],
    {
        "web_search": "web_search",
        "scientific_search": "scientific_search"
    }
)
workflow.add_edge("web_search", "supervisor")
workflow.add_edge("scientific_search", "supervisor")
workflow.add_edge("supervisor", END)

# Compila o grafo para torná-lo executável
app = workflow.compile()

# --- 4. Função para Interface do Gradio ---
def run_graph(user_query: str) -> str:
    """
    Executa o grafo do LangGraph com a pergunta do usuário e retorna a resposta final.
    """
    try:
        result = app.invoke({"user_query": user_query})
        # Retorna a resposta final do supervisor, com um fallback
        return result.get("final_answer", "Não foi possível obter uma resposta final.")
    except Exception as e:
        return f"Ocorreu um erro: {e}"

# --- 5. Lançamento da Aplicação com Gradio ---
if __name__ == "__main__":
    print("Iniciando a interface do Gradio...")
    iface = gr.Interface(
        fn=run_graph,
        inputs=gr.Textbox(lines=2, label="Digite sua pergunta:", placeholder="Ex: Qual o impacto da IA na medicina?"),
        outputs=gr.Markdown(label="Resposta do Agente"),
        title="🤖 Agente de Pesquisa com LangGraph e Gemini",
        description="Faça uma pergunta e obtenha uma resposta com fontes web ou fontes científicas.",
        allow_flagging="never"
    )

    iface.launch()