#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创建Telco电信客户流失分析手册向量数据库
使用现有的专家文档创建FAISS向量数据库
"""

import os
import sys
from dotenv import load_dotenv
load_dotenv(override=True)

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
import pathlib

def create_vector_database():
    """创建向量数据库"""
    print("🔧 开始创建Telco电信客户流失分析手册向量数据库...")
    
    # 初始化embeddings
    embed = OpenAIEmbeddings(
        api_key=os.getenv("OPENAI_API_KEY"),
        base_url="https://ai.devtool.tech/proxy/v1",
        model="text-embedding-3-small",
    )
    
    # 文档路径
    doc_path = "Telco 电信用户流失预测数据集专家文档.md"
    
    if not os.path.exists(doc_path):
        print(f"❌ 错误：找不到文档文件 {doc_path}")
        return False
    
    try:
        # 加载文档
        print("📖 加载专家文档...")
        loader = TextLoader(doc_path, encoding='utf-8')
        documents = loader.load()
        
        # 分割文档
        print("✂️ 分割文档...")
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=["\n\n", "\n", "。", ".", " ", ""]
        )
        
        texts = text_splitter.split_documents(documents)
        print(f"📄 文档已分割为 {len(texts)} 个片段")
        
        # 创建向量数据库
        print("🔍 创建向量数据库...")
        vector_store = FAISS.from_documents(texts, embed)
        
        # 保存向量数据库
        VS_PATH = "telco_customer_churn_analytics_handbook"
        print(f"💾 保存向量数据库到 {VS_PATH}...")
        
        # 确保目录存在
        pathlib.Path(VS_PATH).mkdir(parents=True, exist_ok=True)
        
        vector_store.save_local(VS_PATH)
        
        print("✅ 向量数据库创建成功！")
        print(f"📁 数据库保存位置: {VS_PATH}")
        print(f"📊 包含 {len(texts)} 个文档片段")
        
        return True
        
    except Exception as e:
        print(f"❌ 创建向量数据库时发生错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def test_vector_database():
    """测试向量数据库"""
    print("\n🧪 测试向量数据库...")
    
    try:
        # 初始化embeddings
        embed = OpenAIEmbeddings(
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url="https://ai.devtool.tech/proxy/v1",
            model="text-embedding-3-small",
        )
        
        # 加载向量数据库
        VS_PATH = "telco_customer_churn_analytics_handbook"
        vector_store = FAISS.load_local(
            folder_path=VS_PATH,
            embeddings=embed,
            allow_dangerous_deserialization=True,
        )
        
        # 测试查询
        test_queries = [
            "什么是客户流失？", 
            "影响客户流失的主要因素有哪些？",
            "合同类型如何影响客户流失？",
            "数据集包含哪些字段？"
        ]
        
        for query in test_queries:
            print(f"\n🔍 查询: {query}")
            docs = vector_store.similarity_search(query, k=2)
            if docs:
                print(f"📄 找到 {len(docs)} 个相关文档片段")
                print(f"📝 第一个片段预览: {docs[0].page_content[:100]}...")
            else:
                print("❌ 未找到相关文档")
        
        print("\n✅ 向量数据库测试完成！")
        return True
        
    except Exception as e:
        print(f"❌ 测试向量数据库时发生错误: {str(e)}")
        return False

def main():
    """主函数"""
    print("🚀 Telco电信客户流失分析手册向量数据库创建工具")
    print("=" * 60)
    
    # 检查环境变量
    if not os.getenv("OPENAI_API_KEY"):
        print("❌ 错误：请设置OPENAI_API_KEY环境变量")
        return
    
    # 创建向量数据库
    success = create_vector_database()
    
    if success:
        # 测试向量数据库
        test_vector_database()
        
        print("\n🎉 向量数据库创建和测试完成！")
        print("现在可以正常使用RAG Agent了。")
    else:
        print("\n❌ 向量数据库创建失败！")

if __name__ == "__main__":
    main()