# DESIGN DECISION: Use OllamaEmbeddings and ChatOllama instead of OpenAI to satisfy assignment requirement of local LLM and embeddings via Ollama.
# NECESSITY: Assignment explicitly requires local LLM and embeddings via Ollama; using OpenAI would violate constraints and introduce API keys.
# OPTIMALITY: Ollama provides zero-cost inference, lower latency, and full data control; no external network calls.
# ALTERNATIVES CONSIDERED: OpenRouter or OpenAI; rejected due to requirement of local models and cost.

import os
import sys
import asyncio
from typing import List

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_core.documents import Document
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain_core.messages import HumanMessage
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from qdrant_client import QdrantClient

# Initialize embeddings and chat models
embeddings = OllamaEmbeddings(model="nomic-embed-text")
chat = ChatOllama(model="llama3")

# Initialize Qdrant client and vector store
qdrant_client = QdrantClient(url="http://localhost:6333")
vector_store = QdrantVectorStore(
    client=qdrant_client,
    collection_name="knowledge",
    embedding_function=embeddings,
)

# Text splitter for chunking documents
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

# Tool: Add content to knowledge base
@tool
def add_to_knowledge_base(content: str, title: str = "doc") -> str:
    """Add content to the knowledge base."""
    chunks: List[str] = splitter.split_text(content)
    docs: List[Document] = [
        Document(page_content=chunk, metadata={"title": title}) for chunk in chunks
    ]
    vector_store.add_documents(docs)
    return f"Added {len(docs)} chunks for {title}"

# Tool: Search knowledge base
@tool
def search_knowledge_base(query: str, max_results: int = 3) -> str:
    """Search the knowledge base for relevant information."""
    docs: List[Document] = vector_store.similarity_search(query, k=max_results)
    if not docs:
        return "No results."
    return "\n".join(
        f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)
    )

# Backend for deepagents
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# System prompt guiding the agent
system_prompt = (
    "You are a helpful agent with access to a knowledge base. "
    "Use the provided tools to search and add information. "
    "When searching, return concise results. "
    "When adding, confirm the number of chunks added."
)

# Create the deep agent
agent = create_deep_agent(
    model=chat,
    tools=[add_to_knowledge_base, search_knowledge_base],
    backend=backend,
    system_prompt=system_prompt,
)

# Load documents from a directory into the knowledge base
def load_documents_from_dir(dir_path: str) -> None:
    """Load all .txt files from dir_path into the knowledge base."""
    for root, _, files in os.walk(dir_path):
        for file in files:
            if file.lower().endswith(".txt"):
                path = os.path.join(root, file)
                with open(path, "r", encoding="utf-8") as f:
                    content = f.read()
                title = os.path.splitext(file)[0]
                add_to_knowledge_base(content, title)

# Interactive CLI
async def interactive_loop() -> None:
    print("Welcome to the RAG agent CLI.")
    print("Commands: /add, /search, /quit")
    while True:
        user_input = input("\n> ").strip()
        if user_input.lower() == "/quit":
            print("Goodbye!")
            break
        elif user_input.lower() == "/add":
            title = input("Title: ").strip()
            print("Enter content (end with a single line containing only 'END'):")
            lines: List[str] = []
            while True:
                line = input()
                if line.strip() == "END":
                    break
                lines.append(line)
            content = "\n".join(lines)
            message = f"Add the following content to knowledge base with title '{title}'."
            result = await agent.ainvoke(
                {"messages": [HumanMessage(content=message)]},
                {"configurable": {"thread_id": "session-1"}},
            )
            print(result["messages"][-1].content)
        elif user_input.lower() == "/search":
            query = input("Query: ").strip()
            message = f"Search knowledge base for: {query}"
            result = await agent.ainvoke(
                {"messages": [HumanMessage(content=message)]},
                {"configurable": {"thread_id": "session-1"}},
            )
            print(result["messages"][-1].content)
        else:
            # Treat as normal message
            result = await agent.ainvoke(
                {"messages": [HumanMessage(content=user_input)]},
                {"configurable": {"thread_id": "session-1"}},
            )
            print(result["messages"][-1].content)

async def main() -> None:
    # Optional loading of documents via command line
    if len(sys.argv) > 1 and sys.argv[1] == "--load-dir":
        if len(sys.argv) < 3:
            print("Usage: python main.py --load-dir <directory>")
            return
        dir_path = sys.argv[2]
        if not os.path.isdir(dir_path):
            print(f"Directory not found: {dir_path}")
            return
        print(f"Loading documents from {dir_path}...")
        load_documents_from_dir(dir_path)
        print("Loading complete.")
    await interactive_loop()

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