from phi.vectordb.lancedb import LanceDb
from phi.knowledge.pdf import PDFKnowledgeBase, PDFReader
from phi.embedder.google import GeminiEmbedder
from phi.vectordb.search import SearchType
from phi.utils.log import logger
import os

def load_knowledge_base():
    """
    Loads or creates a knowledge base from PDF documents in the 'knowledge' folder
    using LanceDB for storage.

    This version includes a manual loop to process one file at a time,
    which is a robust workaround for a bug in the library's multi-file handling.
    """
    knowledge_dir = "./knowledge"
    db_dir = "./vectordb/lance_db"
    table_name = "local_pdf_knowledge" 

    if not os.path.exists(knowledge_dir) or not os.listdir(knowledge_dir):
        logger.warning(f"The '{knowledge_dir}' directory is empty or does not exist. No local knowledge base will be loaded.")
        return None

    # Get a list of all PDF files to process
    pdf_files = [f for f in os.listdir(knowledge_dir) if f.lower().endswith(".pdf")]
    if not pdf_files:
        logger.warning(f"No PDF files found in the '{knowledge_dir}' directory.")
        return None

    logger.info("Loading Knowledge Base using LanceDb with manual file processing...")
    
    try:
        knowledge_base = PDFKnowledgeBase(
                path=knowledge_dir, 
                vector_db=LanceDb(
                    table_name=table_name,
                    uri=db_dir,
                    embedder=GeminiEmbedder(model="models/text-embedding-004"),
                    search_type=SearchType.keyword,
                ),
            reader=PDFReader(chunk=True)  
            )
        
        logger.info("All files processed successfully.")

        return knowledge_base

    except Exception as e:
        logger.error(f"An unexpected error occurred during manual file loading: {e}")
        return None