import os

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from huggingface_hub import snapshot_download
import shutil

DATA_FOLDER = "data"
VECTOR_FOLDER = "vector_store"

def download_dataset():

    if os.path.exists(DATA_FOLDER) and os.listdir(DATA_FOLDER):
        print("PDFs already available.")
        return

    print("Downloading PDFs from Hugging Face Dataset...")

    dataset_path = snapshot_download(
        repo_id="LoreSandhu/mediassist-pdfs",
        repo_type="dataset"
    )

    source = dataset_path

    if os.path.isdir(os.path.join(dataset_path, "data")):
        source = os.path.join(dataset_path, "data")

    shutil.copytree(source, DATA_FOLDER, dirs_exist_ok=True)

    print("Dataset downloaded successfully.")

# -----------------------------
# Embedding Model (Load Once)
# -----------------------------

embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)


# -----------------------------
# Load PDFs
# -----------------------------
def load_documents():

    documents = []

    for root, _, files in os.walk(DATA_FOLDER):

        for file in files:

            if file.endswith(".pdf"):

                loader = PyPDFLoader(
                    os.path.join(root, file)
                )

                documents.extend(loader.load())

    return documents

# -----------------------------
# Build Vector Database
# -----------------------------

def build_vector_db():

    download_dataset()
    
    docs = load_documents()

    splitter = RecursiveCharacterTextSplitter(

        chunk_size=500,
        chunk_overlap=80,
        separators=[
            "\n\n",
            "\n",
            ". ",
            " ",
            ""
        ]

    )

    chunks = splitter.split_documents(docs)

    vector_db = FAISS.from_documents(
        chunks,
        embeddings
    )

    vector_db.save_local(VECTOR_FOLDER)

    print(f"\nIndexed {len(chunks)} chunks.")


def load_vector_db():

    if not os.path.exists(os.path.join(VECTOR_FOLDER, "index.faiss")):

        print("Vector database not found. Building...")

        build_vector_db()

    return FAISS.load_local(

        VECTOR_FOLDER,

        embeddings,

        allow_dangerous_deserialization=True

    )

VECTOR_DB = load_vector_db()


# -----------------------------
# Search
# -----------------------------

def search_documents(query, k=6):

    print("\n==============================")
    print("SEARCH QUERY:", query)
    print("==============================")

    retriever = VECTOR_DB.as_retriever(

    search_type="mmr",

    search_kwargs={
        "k": k,
        "fetch_k": 12
    }

)

    docs = retriever.invoke(query)

    if not docs:

        print("No documents retrieved.")
        return ""

    context = ""

    for i, doc in enumerate(docs):

        print(f"\nResult {i+1}")

        print("Source:", doc.metadata)

        print(doc.page_content[:500])

        print("-" * 60)

        context += doc.page_content + "\n\n"

    return context


# -----------------------------
# Build Index
# -----------------------------

if __name__ == "__main__":

    build_vector_db()