"""
RAG agent implementation with Qdrant and Ollama.
"""

import os
from pathlib import Path
from typing import List, Dict

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain.tools import tool
from langchain_core.messages import HumanMessage
from langchain.agents import create_agent

# Initialize LLM and embeddings using Ollama
LLM_MODEL = "llama3"
EMBEDDING_MODEL = "nomic-embed-text"

llm = ChatOllama(model=LLM_MODEL, temperature=0.0)
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

# Qdrant client (in‑memory for simplicity)
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams

client = QdrantClient(":memory:")
COLLECTION_NAME = "knowledge"
if not client.collection_exists(COLLECTION_NAME):
    client.create_collection(
        COLLECTION_NAME,
        vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE),
    )
vector_store = QdrantVectorStore(client=client, collection_name=COLLECTION_NAME, embedding=embeddings)

# Tool: add to knowledge base
@tool
def add_to_knowledge_base(content: str, title: str = "document") -> str:
    """Add a document to the vector store.

    Parameters
    ----------
    content: str
        Raw text of the document.
    title: str, optional
        Title or identifier for the document.
    Returns
    -------
    str
        Confirmation message.
    """
    # Split into chunks using chunker module
    from chunker import split_text
    chunks = split_text(content)
    docs = []
    for i, chunk in enumerate(chunks):
        meta = {"title": title, "chunk_index": str(i)}
        docs.append({"page_content": chunk, "metadata": meta})
    vector_store.add_documents(docs)
    return f"Added {len(chunks)} chunks from '{title}'."

# Tool: search knowledge base
@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Semantic search in the vector store.

    Parameters
    ----------
    query: str
        Search query.
    max_results: int, optional
        Number of top results to return.
    Returns
    -------
    str
        Formatted search results.
    """
    docs = vector_store.similarity_search(query, k=max_results)
    if not docs:
        return "No relevant documents found."
    lines = []
    for i, doc in enumerate(docs, 1):
        title = doc.metadata.get("title", "unknown")
        chunk_idx = doc.metadata.get("chunk_index", "0")
        lines.append(f"{i}. [{title} - chunk {chunk_idx}]\n{doc.page_content[:200]}...")
    return "\n\n".join(lines)

# Create agent with tools
SYSTEM_PROMPT = (
    "You are an assistant that can search and add documents to a knowledge base."
    " Use the provided tools to manage the knowledge base."
)
agent = create_agent(
    llm=llm,
    tools=[add_to_knowledge_base, search_knowledge_base],
    system_prompt=SYSTEM_PROMPT,
)

# Expose agent for external use
__all__ = ["agent", "add_to_knowledge_base", "search_knowledge_base"]