"""
RAG Engine — Handles document parsing, chunking, and ChromaDB vector storage.
"""

import os
import shutil
import uuid
import tempfile
from pathlib import Path

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader, TextLoader, WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# We store the ChromaDB locally
DB_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "chroma_data")
os.makedirs(DB_DIR, exist_ok=True)

_embeddings = None

def get_embeddings():
    global _embeddings
    if not _embeddings:
        _embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    return _embeddings

def get_kb_collection(kb_id: str) -> Chroma:
    """Gets the Chroma vectorstore for a specific knowledge base (collection)."""
    return Chroma(
        collection_name=kb_id,
        embedding_function=get_embeddings(),
        persist_directory=DB_DIR
    )

def list_knowledge_bases():
    """Lists all available knowledge bases by inspecting the Chroma directory or client."""
    # Since Chroma 0.4+, we interact with the persistent client directly to list collections
    import chromadb
    client = chromadb.PersistentClient(path=DB_DIR)
    collections = client.list_collections()
    
    # Each collection is a kb. Return list of dicts.
    return [{"id": c.name, "name": c.name} for c in collections]

def ingest_document(kb_id: str, file_path_or_url: str, doc_type: str):
    """
    Ingests a document or URL into the specified knowledge base.
    """
    if doc_type == "pdf":
        loader = PyPDFLoader(file_path_or_url)
        docs = loader.load()
    elif doc_type == "txt":
        loader = TextLoader(file_path_or_url, encoding="utf-8")
        docs = loader.load()
    elif doc_type == "json":
        import json
        from langchain_core.documents import Document
        with open(file_path_or_url, "r", encoding="utf-8") as f:
            try:
                data = json.load(f)
                text_content = json.dumps(data, indent=2)
            except json.JSONDecodeError:
                f.seek(0)
                text_content = f.read()
        docs = [Document(page_content=text_content, metadata={"source": file_path_or_url})]
    elif doc_type == "docx":
        from langchain_community.document_loaders import Docx2txtLoader
        loader = Docx2txtLoader(file_path_or_url)
        docs = loader.load()
    elif doc_type == "url":
        loader = WebBaseLoader(file_path_or_url)
        docs = loader.load()
    else:
        raise ValueError(f"Unsupported document type: {doc_type}")
    
    # Chunking
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", " ", ""]
    )
    splits = text_splitter.split_documents(docs)
    
    # Store in ChromaDB
    vectorstore = get_kb_collection(kb_id)
    vectorstore.add_documents(documents=splits)
    
    return len(splits)

def search_kb(kb_id: str, query: str, top_k: int = 3) -> str:
    """
    Searches the knowledge base and returns a formatted string of the top context chunks.
    """
    try:
        vectorstore = get_kb_collection(kb_id)
        docs = vectorstore.similarity_search(query, k=top_k)
        if not docs:
            return "No relevant information found in the knowledge base."
        
        results = []
        for i, doc in enumerate(docs):
            source = doc.metadata.get("source", "Unknown Source")
            page = doc.metadata.get("page", "")
            page_info = f" (Page {page})" if page else ""
            results.append(f"--- Doc {i+1} | Source: {source}{page_info} ---\n{doc.page_content}")
            
        return "\n\n".join(results)
    except Exception as e:
        print(f"Error searching KB {kb_id}: {e}")
        return f"Error retrieving from Knowledge Base: {e}"
