"""
Main agent implementation using LangChain.
"""

import os
from typing import Dict

from langchain import LLMChain, PromptTemplate
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import BaseTool

from rag_tools import search_local_kb, web_search
from vectorstore import create_vectorstore, load_documents

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_NAME = "llama3"
CHROMA_DIR = "./chroma_db"
DOCS_DIR = "documents"

# ---------------------------------------------------------------------------
# Load or create vector store
# ---------------------------------------------------------------------------
vectorstore = create_vectorstore(CHROMA_DIR)
# Load documents only if the store is empty
if not vectorstore._collection.count():
    load_documents(DOCS_DIR, vectorstore)

# ---------------------------------------------------------------------------
# Define tools
# ---------------------------------------------------------------------------
class LocalKBTool(BaseTool):
    name = "search_local_kb"
    description = "Perform a semantic search in the local knowledge base."

    def _run(self, query: str, top_k: int = 3) -> str:  # pragma: no cover
        return search_local_kb(query, top_k, vectorstore)

class WebSearchTool(BaseTool):
    name = "web_search"
    description = "Search the web using Tavily."

    def _run(self, query: str, top_k: int = 3) -> str:  # pragma: no cover
        return web_search(query, top_k)

tools = [LocalKBTool(), WebSearchTool()]

# ---------------------------------------------------------------------------
# Prompt template
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """
You are an AI assistant that answers user questions.
- If the answer can be found in the local knowledge base, use the tool `search_local_kb`.
- If the answer requires up‑to‑date information, use the tool `web_search`.
After providing the answer, always state the source in the format:

Source: <chromadb|tavily>
"""

PROMPT = PromptTemplate(
    input_variables=["input", "chat_history"],
    template="""
{chat_history}
User: {input}
Assistant: """
)

# ---------------------------------------------------------------------------
# Agent chain
# ---------------------------------------------------------------------------
llm = ChatOllama(model=MODEL_NAME, temperature=0.2)
chain = LLMChain(llm=llm, prompt=PROMPT)

# ---------------------------------------------------------------------------
# Helper to decide which tool to use
# ---------------------------------------------------------------------------
def decide_and_run(query: str) -> Dict[str, str]:
    """Use the LLM to decide whether to use local KB or web search.

    Returns a dict with keys: answer, source.
    """
    # Simple heuristic: if the query contains words like "news", "latest", "today" use web
    web_keywords = {"news", "latest", "today", "current", "recent", "update"}
    if any(word in query.lower() for word in web_keywords):
        result = web_search(query)
        source = "tavily"
    else:
        result = search_local_kb(query, vectorstore=vectorstore)
        source = "chromadb"
    return {"answer": result, "source": source}

# ---------------------------------------------------------------------------
# CLI loop
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    print("RAG Agent ready. Type 'exit' to quit.")
    while True:
        user_input = input("\nЗапрос: ")
        if user_input.lower() in {"exit", "quit", "q"}:
            break
        output = decide_and_run(user_input)
        print(f"\nОтвет:\n{output['answer']}")
        print(f"Источник: {output['source']}")