import os
from pathlib import Path
from dotenv import load_dotenv

from langchain_ollama import ChatOllama
from langchain.agents import Tool, AgentExecutor, create_openai_tools_agent
from langchain.tools import tool
from langchain.schema import HumanMessage

from vectorstore import create_vectorstore, load_documents

# Load env
load_dotenv()
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
if not TAVILY_API_KEY:
    raise RuntimeError("TAVILY_API_KEY not set in .env")

# Setup vector store
vectorstore = create_vectorstore()
# Load documents if not already loaded
if not Path("./chroma_db/chroma-collections.jsonl").exists():
    load_documents("documents", vectorstore)

retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

@tool(name="search_local_kb", description="Search local knowledge base in ChromaDB")
def search_local_kb(query: str, top_k: int = 3) -> str:
    docs = retriever.invoke({"query": query, "k": top_k})
    return "\n---\n".join([d.page_content for d in docs])

@tool(name="web_search", description="Search the web via Tavily")
def web_search(query: str) -> str:
    from tavily import TavilyClient
    client = TavilyClient(api_key=TAVILY_API_KEY)
    results = client.search(query, max_results=3)
    return "\n---\n".join([f"{r.title}\n{r.url}" for r in results])

tools = [search_local_kb, web_search]

system_prompt = (
    "You are an AI assistant that answers user questions.
    If the answer can be found in local documents, use search_local_kb.\n"
    "If the question is about recent events or requires up-to-date info, use web_search.\n"
    "Always indicate the source of your answer: chromadb or tavily."
)

agent = create_openai_tools_agent(
    llm=ChatOllama(model="llama3", temperature=0),
    tools=tools,
    system_message=system_prompt,
)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

print("RAG agent ready. Type 'exit' to quit.")
while True:
    user_input = input("Query: ")
    if user_input.lower() in {"exit", "quit"}:
        break
    response = executor.invoke({"input": user_input})
    print(response["output"])
print("Goodbye!")