import_data.py

import os
import weaviate
from llama_index import download_loader
from llama_index.vector_stores import WeaviateVectorStore
from llama_index import VectorStoreIndex, StorageContext
from pathlib import Path
import argparse


def get_pdf_files(base_path, loader):
    """
    Get paths to all PDF files in a directory and its subdirectories.

    Parameters:
    - base_path (str): The path to the starting directory.

    Returns:
    - list of str: A list of paths to all PDF files found.
    """
    pdf_paths = []

    # Check if the base path exists and is a directory
    if not os.path.exists(base_path):
        raise FileNotFoundError(f"The specified base path does not exist: {base_path}")
    if not os.path.isdir(base_path):
        raise NotADirectoryError(
            f"The specified base_path is not a directory: {base_path}"
        )

    # Loop through all directories and files starting from the base path
    for root, dirs, files in os.walk(base_path):
        for filename in files:
            # If a file has a .pdf extension, add its path to the list
            if filename.endswith(".pdf"):
                pdf_file = loader.load_data(file=Path(root, filename))
                pdf_paths.extend(pdf_file)

    return pdf_paths


def main(args):
    PDFReader = download_loader("PDFReader")
    loader = PDFReader()

    documents = get_pdf_files(args.pdf_dir, loader)

    client = weaviate.Client(
        url=os.environ["WEAVIATE_URL"],
        auth_client_secret=weaviate.AuthApiKey(api_key=os.environ["WEAVIATE_API_KEY"]),
        additional_headers={"X-OpenAI-Api-Key": os.environ["OPENAI_API_KEY"]},
    )

    # construct vector store
    vector_store = WeaviateVectorStore(
        weaviate_client=client, index_name=args.customer, text_key="content"
    )

    # setting up the storage for the embeddings
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    # set up the index
    index = VectorStoreIndex(documents, storage_context=storage_context)
    query_engine = index.as_query_engine()
    response = query_engine.query(args.query)
    print(response)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process and query PDF files.")

    parser.add_argument("--customer", default="Ausy", help="Customer name")
    parser.add_argument("--pdf_dir", default="./data", help="Directory containing PDFs")
    parser.add_argument(
        "--query",
        default="What is CX0 customer exprience office?",
        help="Query to execute",
    )

    args = parser.parse_args()

    main(args)