"""LangGraph Agent"""
import os
from dotenv import load_dotenv
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from langchain_community.vectorstores import SupabaseVectorStore
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from supabase import create_client, Client

# Cargar variables de entorno
load_dotenv()

# Configuración de Supabase (segura)
try:
    supabase_url = os.environ['SUPABASE_URL']
    supabase_key = os.environ['SUPABASE_KEY']
    supabase: Client = create_client(supabase_url, supabase_key)
except KeyError as e:
    raise ValueError(f"""
    Missing Supabase credentials. Please set:
    1. SUPABASE_URL and SUPABASE_KEY as Secrets in HF Space settings
    OR
    2. As environment variables in a .env file for local development
    Missing: {e}
    """)

# Herramientas (tools) permanecen igual...
# [Mantén todas tus herramientas (@tool) como están]

# Cargar prompt del sistema
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()

sys_msg = SystemMessage(content=system_prompt)

# Configuración del vector store corregida
try:
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    vector_store = SupabaseVectorStore(
        client=supabase,
        embedding=embeddings,
        table_name="documents",
        query_name="match_documents_langchain",
    )
    retriever_tool = create_retriever_tool(
        retriever=vector_store.as_retriever(),
        name="Question Search",
        description="Search for similar questions in our knowledge base",
    )
    tools.append(retriever_tool)  # Añade esta herramienta a tu lista existente
except Exception as e:
    print(f"Warning: Could not initialize vector store: {e}")
    retriever_tool = None

# Función build_graph corregida
def build_graph(provider: str = "groq"):
    """Build the graph with error handling"""
    try:
        if provider == "google":
            llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
        elif provider == "groq":
            llm = ChatGroq(model="mixtral-8x7b-32768", temperature=0)
        elif provider == "huggingface":
            llm = ChatHuggingFace(
                llm=HuggingFaceEndpoint(
                    endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2",
                    temperature=0,
                )
            )
        else:
            raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'")
            
        llm_with_tools = llm.bind_tools(tools)
        
        # Nodos del grafo
        def assistant(state: MessagesState):
            return {"messages": [llm_with_tools.invoke(state["messages"])]}
            
        def retriever(state: MessagesState):
            if not retriever_tool:
                return {"messages": state["messages"]}
            similar_questions = vector_store.similarity_search(state["messages"][-1].content, k=1)
            if similar_questions:
                example_msg = HumanMessage(
                    content=f"Similar question found:\n\n{similar_questions[0].page_content}"
                )
                return {"messages": state["messages"] + [example_msg]}
            return {"messages": state["messages"]}
            
        # Construcción del grafo
        builder = StateGraph(MessagesState)
        builder.add_node("retriever", retriever)
        builder.add_node("assistant", assistant)
        builder.add_node("tools", ToolNode(tools))
        
        builder.add_edge(START, "retriever")
        builder.add_edge("retriever", "assistant")
        builder.add_conditional_edges(
            "assistant",
            tools_condition,
        )
        builder.add_edge("tools", "assistant")
        
        return builder.compile()
        
    except Exception as e:
        print(f"Error building graph: {e}")
        raise

# [El resto del código permanece igual]