"""Main module implementing a RAG-enabled agent with Qdrant and Ollama.

The code follows the assignment specification and uses only the required libraries.
"""

import os
import sys
import json
from pathlib import Path
from typing import List, Dict, Any

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor
from langchain.schema import AgentAction, AgentFinish
from langchain_core.prompts import ChatPromptTemplate

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_COLLECTION = "knowledge"
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3"
CHUNK_SIZE = 1000  # characters
CHUNK_OVERLAP = 200

# ---------------------------------------------------------------------------
# Vector store helper
# ---------------------------------------------------------------------------
class KnowledgeBase:
    """Wrapper around QdrantVectorStore providing add/search helpers."""

    def __init__(self, url: str = QDRANT_URL, collection: str = QDRANT_COLLECTION):
        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
        self.store = QdrantVectorStore.from_existing_collection(
            collection_name=collection,
            url=url,
            embedding=self.embeddings,
        )

    def add_documents(self, documents: List[str], titles: List[str]):
        """Adds a list of documents with corresponding titles to the store.

        Each document is split into chunks before being stored.
        """
        splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
        for doc, title in zip(documents, titles):
            chunks = splitter.split_text(doc)
            metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))]
            self.store.add_texts(chunks, metadatas=metadatas)

    def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
        """Semantic search in the knowledge base.

        Returns a list of dicts with keys: text, title, distance.
        """
        results = self.store.similarity_search_with_score(query, k=max_results)
        return [
            {
                "text": text,
                "title": meta.get("title", "unknown"),
                "distance": score,
            }
            for text, score, meta in results
        ]

# ---------------------------------------------------------------------------
# Global knowledge base instance
# ---------------------------------------------------------------------------
kb = KnowledgeBase()

# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------
@tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Search the knowledge base for a query and return formatted results."""
    results = kb.search(query, max_results)
    if not results:
        return "No relevant documents found."
    formatted = [f"Title: {r['title']}\nSnippet: {r['text'][:200]}...\nDistance: {r['distance']:.4f}" for r in results]
    return "\n\n".join(formatted)

@tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str:
    """Add a new document to the knowledge base."""
    kb.add_documents([content], [title])
    return f"Document '{title}' added successfully."

# ---------------------------------------------------------------------------
# Agent definition
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = (
    "You are an assistant that can search and add documents to a local knowledge base. "
    "Use the provided tools to perform semantic search and store new information. "
    "When answering user queries, first determine if the user needs a search or an addition. "
    "If no relevant information is found, suggest adding new content."
)

prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("user", "{input}"),
])

agent = create_agent(
    llm=ChatOllama(model=LLM_MODEL),
    tools=[search_knowledge_base, add_to_knowledge_base],
    prompt=prompt,
    verbose=True,
)

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

# ---------------------------------------------------------------------------
# CLI client
# ---------------------------------------------------------------------------
def load_documents_from_dir(directory: str):
    """Load all .txt files from a directory and add them to the knowledge base."""
    docs = []
    titles = []
    for path in Path(directory).glob("**/*.txt"):
        text = path.read_text(encoding="utf-8")
        docs.append(text)
        titles.append(path.stem)
    if docs:
        kb.add_documents(docs, titles)
        print(f"Loaded {len(docs)} documents from {directory}.")
    else:
        print("No .txt files found.")


def interactive_loop():
    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /load <dir>, /quit")
    while True:
        try:
            user_input = input("> ")
        except (KeyboardInterrupt, EOFError):
            print("\nExiting.")
            break
        if not user_input:
            continue
        if user_input.startswith("/quit"):
            print("Goodbye!")
            break
        elif user_input.startswith("/add "):
            try:
                _, rest = user_input.split("/add ", 1)
                title, content = rest.split(" ", 1)
                response = add_to_knowledge_base(content, title)
                print(response)
            except ValueError:
                print("Usage: /add <title> <content>")
        elif user_input.startswith("/search "):
            query = user_input.split("/search ", 1)[1]
            results = search_knowledge_base(query)
            print(results)
        elif user_input.startswith("/load "):
            dir_path = user_input.split("/load ", 1)[1]
            load_documents_from_dir(dir_path)
        else:
            # Treat as normal user query
            result = agent_executor.invoke({"input": user_input})
            print(result.get("output", ""))

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "cli":
        interactive_loop()
    else:
        print("Usage: python main.py cli")
        print("Run the interactive CLI with: python main.py cli")