import os
from pathlib import Path

from dotenv import load_dotenv

# LlamaIndex 核心组件
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage,
    Settings,
)
from llama_index.embeddings.openai import OpenAIEmbedding

# ============================
# 环境变量 & 路径配置
# ============================

# 加载 .env（和 Clare 项目保持一致，直接复用 OPENAI_API_KEY）
load_dotenv()

PROJECT_ROOT = Path(__file__).resolve().parent

# 1. GENAI 课程目录（你要向量化的本地课程代码 / 笔记）
#   这里用绝对路径更稳，不怕你从哪里运行脚本
EXPORT_DIR = PROJECT_ROOT / "GENAI COURSES"

# 2. 向量数据库持久化路径
PERSIST_DIR = PROJECT_ROOT / "genai_courses_index"

# 3. 显式指定 Embedding 模型（和 Clare 一致：text-embedding-3-small）
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

def get_index(force_rebuild=False):
    """
    获取索引：优先读取本地缓存，如果不存在或强制刷新，则重新构建
    """
    
    # 情况 A: 数据库已存在，直接加载 (极快，不花钱)
    if PERSIST_DIR.exists() and not force_rebuild:
        print(f"📂 发现本地数据库 ({PERSIST_DIR})，正在加载...")
        try:
            storage_context = StorageContext.from_defaults(persist_dir=str(PERSIST_DIR))
            index = load_index_from_storage(storage_context)
            print("✅ 加载成功！")
            return index
        except Exception as e:
            print(f"⚠️ 本地数据库加载失败: {e}，准备重新构建...")

    # 情况 B: 数据库不存在，或者要求强制更新 -> 读取本地文件构建
    print(f"🚀 开始扫描本地文件并构建向量数据库：{EXPORT_DIR}")
    if not EXPORT_DIR.exists():
        raise FileNotFoundError(f"GENAI COURSES 目录不存在：{EXPORT_DIR}")
    
    # 1. 读取文件
    # recursive=True 会读取子文件夹，确保附件和嵌套页面都被读取
    # required_exts 可以指定只读 .md，如果不加这行则会读取所有支持的文件(pdf, txt, etc.)
    reader = SimpleDirectoryReader(
        input_dir=str(EXPORT_DIR),
        recursive=True,
        # 可以根据需要调整：这里把常见的课程文件类型都包含进来
        required_exts=[".md", ".pdf", ".txt", ".py", ".ipynb"],
    )
    documents = reader.load_data()
    print(f"📄 成功读取了 {len(documents)} 个文件片段")

    # 2. 构建索引 (这一步会调用 OpenAI API 进行 Embedding)
    print("🧠 正在生成向量索引 (Embedding)...")
    index = VectorStoreIndex.from_documents(documents)

    # 3. 保存到硬盘
    print(f"💾 正在保存数据库到 {PERSIST_DIR} ...")
    index.storage_context.persist(persist_dir=PERSIST_DIR)
    
    return index

if __name__ == "__main__":
    # --- 主程序 ---
    
    # 首次运行或通过参数控制 force_rebuild=True 来更新
    index = get_index(force_rebuild=False)
    
    # 创建查询引擎
    query_engine = index.as_query_engine()
    
    print("\n💬 本地知识库助手已就绪 (输入 'exit' 退出, 'update' 重建):")
    
    while True:
        question = input("\n请输入问题: ")
        
        if question.lower() == 'exit':
            break
        elif question.lower() == 'update':
            index = get_index(force_rebuild=True)
            query_engine = index.as_query_engine()
            print("✅ 数据库已更新！")
            continue
            
        response = query_engine.query(question)
        print(f"\n🤖 回答:\n{response}")