"""
Agent definition for the RAG system.

Provides two tools:
* search_knowledge_base(query, max_results)
* add_to_knowledge_base(content, title)

The agent is created with create_agent from langchain.agents.
"""

import os
from typing import List
from langchain_ollama import OllamaEmbeddings
from chromadb import PersistentClient
from chromadb.utils import embedding_functions as ef
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage

# Initialize embeddings and Chroma client
EMBEDDINGS = OllamaEmbeddings(model="nomic-embed-text")
CHROMA_PATH = os.path.join(os.getcwd(), "chromadb_store")
CLIENT = PersistentClient(path=CHROMA_PATH)
COLLECTION_NAME = "knowledge"
if COLLECTION_NAME not in CLIENT.list_collections():
    CLIENT.create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
COLL = CLIENT.get_or_create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)

# Chunker from chunker.py
from chunker import CHUNKER

@tool
def add_to_knowledge_base(content: str, title: str) -> str:
    """
    Add a document to the knowledge base.
    The content is split into chunks and stored with metadata.
    Returns confirmation message.
    """
    # Split content
    chunks = CHUNKER.split_text(content)
    ids = [f"{title}_{i}" for i in range(len(chunks))]
    metadatas = [{"title": title} for _ in chunks]
    COLL.add(ids=ids, documents=chunks, metadatas=metadatas)
    return f"Added {len(chunks)} chunks from '{title}'."

@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """
    Search the knowledge base for relevant documents.
    Returns a formatted string of results.
    """
    results = COLL.query(
        query_texts=[query],
        n_results=max_results,
        include=['documents', 'distances'],
    )
    docs = results.get("documents", [])[0]
    dists = results.get("distances", [])[0]
    if not docs:
        return "No relevant documents found."
    output_lines = []
    for i, (doc, dist) in enumerate(zip(docs, dists), 1):
        output_lines.append(f"{i}. (score: {dist:.4f})\n{doc[:200]}...")
    return "\n\n".join(output_lines)

# Create agent
SYSTEM_PROMPT = (
    "You are an assistant that can search and add to a knowledge base. Use the provided tools."
)
AGENT = create_agent(
    llm=None,  # No LLM needed for tool calls; agent will use system prompt only
    tools=[add_to_knowledge_base, search_knowledge_base],
    system_prompt=SYSTEM_PROMPT,
)

# Expose a simple invoke function
async def run_agent(messages: List[HumanMessage]):
    return await AGENT.ainvoke({"messages": messages}, {"configurable": {"thread_id": "rag-agent"}})