from pathlib import Path
from typing import List

from agno.document import Document
from agno.utils.log import logger
from dotenv import load_dotenv

from agentic_rag import initialize_knowledge_base, get_reader, initialize_mingci_knowledge_base

# 加载.env文件
load_dotenv()
import os

def main():
    print("Hello from agno-agentic-rag!")
    # 从.env加载知识库来源目录并初始化知识库
    mingci_knowledge_source_dir = os.getenv("MINGCI_KNOWLEDGE_SOURCE_DIR")
    if mingci_knowledge_source_dir and os.path.exists(mingci_knowledge_source_dir):
        # 初始化知识库
        knowledge_base = initialize_mingci_knowledge_base()

        LoadKnowledgeToDatabase(knowledge_base, mingci_knowledge_source_dir)

    knowledge_source_dir = os.getenv("KNOWLEDGE_SOURCE_DIR")
    if knowledge_source_dir and os.path.exists(knowledge_source_dir):
        # 初始化知识库
        knowledge_base = initialize_knowledge_base()

        LoadKnowledgeToDatabase(knowledge_base, knowledge_source_dir)

def LoadKnowledgeToDatabase(knowledge_base, knowledge_source_dir):
    logger.info(f"加载知识库: {knowledge_source_dir}")
    for root, _, files in os.walk(knowledge_source_dir):
        for file in files:
            file_path = os.path.join(root, file)
            file_ext = os.path.splitext(file)[1][1:]  # 获取文件扩展名
            reader = get_reader(file_ext)
            if reader:
                try:
                    filePath = Path(file_path)
                    docs: List[Document] = reader.read(filePath)
                    knowledge_base.load_documents(docs, upsert=True)
                except Exception as e:
                    logger.warning(f"无法加载文档 {file_path}: {str(e)}")


if __name__ == "__main__":
    main()