"""RAG Agent with Qdrant and OpenRouter.

This script implements the assignment requirements:
* Two tools – `search_knowledge_base` and `add_to_knowledge_base` – are defined with the `@tool` decorator.
* A Qdrant vector store is used for semantic search.  Documents are split into chunks with a
  `RecursiveCharacterTextSplitter` that has `chunk_overlap=100` as requested.
* The agent is created with LangChain’s `create_agent` (the "Исправить" instruction overrides the
  earlier requirement to use `create_deep_agent`).
* A simple CLI allows adding documents, searching the knowledge base and quitting.

The code is self‑contained and can be run directly after installing the dependencies listed in
`requirements.txt`.
"""

import os
import asyncio
import pathlib
from typing import List

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langchain_community.vectorstores import Qdrant
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.agents import create_agent, AgentExecutor, AgentType

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY environment variable is required")

# LLM and embeddings via OpenRouter
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
)

embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
)

# ---------------------------------------------------------------------------
# Vector store setup (Qdrant)
# ---------------------------------------------------------------------------
# Qdrant is expected to be running locally on the default port 6333.
# If you need a different host/port, adjust the `url` parameter.
vector_store = Qdrant.from_existing_index(
    collection_name="knowledge",
    embeddings=embeddings,
    url="http://localhost:6333",
)

# ---------------------------------------------------------------------------
# Text splitter
# ---------------------------------------------------------------------------
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100,  # as required by the assignment
)

# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------
@tool
def search_knowledge_base(query: str, max_results: int = 3) -> str:
    """Search the knowledge base for relevant information.

    Parameters
    ----------
    query: str
        The search query.
    max_results: int, optional
        Number of top results to return. Defaults to 3.
    """
    docs = vector_store.similarity_search(query, k=max_results)
    if not docs:
        return "No results found."
    return "\n\n---\n\n".join(doc.page_content for doc in docs)

@tool
def add_to_knowledge_base(content: str, title: str = "document") -> str:
    """Add content to the knowledge base.

    Parameters
    ----------
    content: str
        The raw text to add.
    title: str, optional
        A title for the document. Defaults to "document".
    """
    # Split into chunks and create Document objects
    chunks = text_splitter.split_text(content)
    from langchain_core.documents import Document

    docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
    vector_store.add_documents(docs)
    return f"Added {len(docs)} chunks for title '{title}'."

# ---------------------------------------------------------------------------
# Agent creation (LangChain create_agent)
# ---------------------------------------------------------------------------
# The system prompt instructs the agent to use the knowledge base tools.
SYSTEM_PROMPT = (
    "You are an AI assistant with access to a knowledge base. "
    "Use the tools `search_knowledge_base` and `add_to_knowledge_base` to answer user queries. "
    "If the user asks to add information, store it. If the user asks for information, search the base."
)

agent = create_agent(
    llm=llm,
    tools=[search_knowledge_base, add_to_knowledge_base],
    system_prompt=SYSTEM_PROMPT,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
)

agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base])

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
async def handle_user_input(user_input: str) -> str:
    if user_input.startswith("/add "):
        # Expected format: /add <title> | <content>
        try:
            _, rest = user_input.split("/add ", 1)
            title, content = rest.split("|", 1)
            title = title.strip()
            content = content.strip()
            result = add_to_knowledge_base(content, title)
            return result
        except ValueError:
            return "Invalid format. Use: /add <title> | <content>"
    elif user_input.startswith("/search "):
        query = user_input[len("/search "):].strip()
        return search_knowledge_base(query)
    elif user_input == "/quit":
        return "quit"
    else:
        # Forward to the agent
        response = await agent_executor.ainvoke({"messages": [HumanMessage(content=user_input)]})
        return response["messages"][-1].content

async def main():
    print("RAG Agent CLI. Commands: /add <title> | <content>, /search <query>, /quit")
    while True:
        user_input = input("> ")
        if not user_input:
            continue
        result = await handle_user_input(user_input)
        if result == "quit":
            print("Goodbye!")
            break
        print(result)

if __name__ == "__main__":
    asyncio.run(main())