#!/usr/bin/env python3
"""RAG‑agent with ChromaDB and Tavily search.

This implementation follows the course specification and uses the
`deepagents` framework to create a single agent that can decide whether
to query the local knowledge base (ChromaDB) or perform a web search
via Tavily.  The agent is backed by OpenRouter for both LLM and
embeddings, complying with the mandatory technical constraints.
"""

import asyncio
import os
from pathlib import Path

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_tavily import TavilySearchRun
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
# Load environment variables (OpenRouter key, Tavily key)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")

# Persist directory for ChromaDB
CHROMA_DIR = Path("./chroma_db")
CHROMA_DIR.mkdir(parents=True, exist_ok=True)

# ---------------------------------------------------------------------------
# 1. Vector store (ChromaDB + OpenRouter embeddings)
# ---------------------------------------------------------------------------
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
)

vector_store = Chroma(
    collection_name="knowledge",
    embedding_function=embeddings,
    persist_directory=str(CHROMA_DIR),
)

# Helper to load documents from a directory and add to the store

def load_documents(directory: Path):
    """Read .txt/.md files, split into chunks, and store in Chroma."""
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs = []
    for file in directory.glob("**/*"):
        if file.suffix.lower() not in {".txt", ".md"}:
            continue
        text = file.read_text(encoding="utf-8")
        chunks = splitter.split_text(text)
        docs.extend([Document(page_content=c, metadata={"source": str(file)}) for c in chunks])
    if docs:
        vector_store.add_documents(docs)
        vector_store.persist()

# ---------------------------------------------------------------------------
# 2. Tools
# ---------------------------------------------------------------------------
@tool
def search_local_kb(query: str, top_k: int = 3) -> str:
    """Semantic search in the local ChromaDB knowledge base."""
    docs = vector_store.similarity_search(query, k=top_k)
    if not docs:
        return "No relevant information found in the local knowledge base."
    return "\n---\n".join([f"{i+1}. {d.page_content[:200]}…" for i, d in enumerate(docs)])

@tool
def web_search(query: str) -> str:
    """Perform a web search using Tavily."""
    tavily = TavilySearchRun(api_key=TAVILY_API_KEY, max_results=3)
    results = tavily.run(query)
    if not results:
        return "No web results found."
    return "\n---\n".join([f"{i+1}. {r['title']}\n{r['content'][:200]}…" for i, r in enumerate(results)])

# ---------------------------------------------------------------------------
# 3. Agent (deepagents)
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
)

backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

system_prompt = (
    "You are a helpful assistant. For a user query, first decide whether the
    answer can be found in the local knowledge base. If so, use the
    `search_local_kb` tool. If the query requires up‑to‑date information,
    use the `web_search` tool. Respond with the best answer and clearly
    state the source: `chromadb` or `tavily`."
)

agent = create_deep_agent(
    model=llm,
    tools=[search_local_kb, web_search],
    backend=backend,
    system_prompt=system_prompt,
)

# ---------------------------------------------------------------------------
# 4. CLI loop
# ---------------------------------------------------------------------------
async def main():
    print("RAG Agent ready. Type your question (or 'exit' to quit).")
    thread_id = "session-1"
    while True:
        user_input = input("\nЗапрос: ")
        if user_input.lower() in {"exit", "quit", "q"}:
            print("Goodbye!")
            break
        response = await agent.ainvoke(
            {"messages": [{"role": "user", "content": user_input}]},
            {"configurable": {"thread_id": thread_id}},
        )
        # The last message contains the assistant reply
        assistant_msg = response["messages"][-1].content
        print(f"\nОтвет:\n{assistant_msg}")

if __name__ == "__main__":
    # Load documents once at startup
    docs_dir = Path("./documents")
    if docs_dir.exists():
        load_documents(docs_dir)
    asyncio.run(main())