"""Main agent logic.

Creates a Chroma vector store, loads documents from the ``documents`` directory,
and runs a simple chat loop. The agent decides whether to use the local KB
or perform a web search based on the presence of the word "news" or
"latest" in the query.
"""

import os
from typing import List

from langchain_ollama import ChatOllama
from langchain.agents import tool, AgentExecutor, ZeroShotAgent
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate

from vectorstore import create_vectorstore, load_documents
from rag_tools import web_search

# ---------------------------------------------------------------------------
# 1. Setup vector store
# ---------------------------------------------------------------------------
VECTORSTORE_DIR = "./chroma_db"
DOCS_DIR = "./documents"

vectorstore = create_vectorstore(VECTORSTORE_DIR)
load_documents(DOCS_DIR, vectorstore)

# ---------------------------------------------------------------------------
# 2. Define local search tool (needs the vectorstore)
# ---------------------------------------------------------------------------
@tool
def search_local_kb(query: str, top_k: int = 3) -> str:
    """Semantic search over the local ChromaDB collection.

    Returns a formatted string containing the top_k snippets and a source tag.
    """
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke(query)
    snippets = "\n".join([f"{idx+1}. {doc.page_content[:200]}" for idx, doc in enumerate(docs)])
    return f"[Local KB]\n{snippets}\nSource: chromadb"

# ---------------------------------------------------------------------------
# 3. Agent prompt and execution
# ---------------------------------------------------------------------------
# The agent will be given two tools: search_local_kb and web_search.
# We provide a simple instruction to choose the appropriate tool.

agent_prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template(
            "You are an assistant that can search a local knowledge base or the web. "
            "If the question is about recent events or news, use web_search. "
            "Otherwise, use search_local_kb. "
            "Respond with the answer and the source (chromadb or tavily)."
        ),
    ]
)

# Create the agent with the two tools
tools = [search_local_kb, web_search]

agent = ZeroShotAgent.from_llm_and_tools(
    llm=ChatOllama(model="llama3", temperature=0),
    tools=tools,
    prompt=agent_prompt,
)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# ---------------------------------------------------------------------------
# 4. Chat loop
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    print("RAG Agent ready. Type 'exit' to quit.")
    while True:
        user_input = input("Запрос: ")
        if user_input.lower() in {"exit", "quit", "q"}:
            print("Bye!")
            break
        try:
            result = agent_executor.invoke({"input": user_input})
            print(result["output"])
        except Exception as e:
            print(f"Error: {e}")