import warnings

from Config import Config
from Config.Config import *
from elasticsearch import Elasticsearch
from gensim.models import KeyedVectors
import jieba
import os
import time

# 抑制HTTPS相关警告
warnings.filterwarnings('ignore', message='Connecting to .* using TLS with verify_certs=False is insecure')
warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host')

# 1. 加载预训练的 Word2Vec 模型
model_path = MODEL_PATH
model = KeyedVectors.load_word2vec_format(model_path, binary=False, limit=MODEL_LIMIT)
print(f"模型加载成功，词向量维度: {model.vector_size}")


# 功能：将文本转换为嵌入向量
def text_to_embedding(text):
    words = jieba.lcut(text)
    embeddings = [model[word] for word in words if word in model]
    if embeddings:
        return sum(embeddings) / len(embeddings)
    return [0.0] * model.vector_size


# 2. 初始化Elasticsearch连接
es = Elasticsearch(
    hosts=Config.ES_CONFIG['hosts'],
    basic_auth=Config.ES_CONFIG['basic_auth'],
    verify_certs=False
)

# 3. 处理processed_chunks目录下的所有文件
txt_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'Txt')

for filename in os.listdir(txt_dir):
    if filename.endswith('.txt'):
        filepath = os.path.join(txt_dir, filename)
        with open(filepath, 'r', encoding='utf-8') as f:
            # 只读取第一行作为向量计算
            first_line = f.readline().strip()
            # 读取全部内容用于后续查询
            full_content = first_line + '\n' + f.read()

            if not first_line:
                print(f"跳过空文件: {filename}")
                continue

            print(f"正在处理文件: {filename}")

            # 4. 获取当前时间和会话ID
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
            # 需要进行标记的标签
            x = filename.split("_")
            selectedTags = [x[0] + "_" + x[1]]
            tags = {"tags": selectedTags, "full_content": full_content}  # 添加完整内容

            # 5. 将第一行文本转换为嵌入向量
            embedding = text_to_embedding(first_line)

            # 6. 插入数据到Elasticsearch
            doc = {
                'tags': tags,
                'user_input': first_line,
                'timestamp': timestamp,
                'embedding': embedding
            }
            es.index(index=ES_CONFIG['index_name'], document=doc)
            print(f"文件 {filename} 数据插入成功")

print("所有文件处理完成")