# agent/embeddings/embedding_utils.py
import json
import numpy as np
from config import *
from zhipuai import ZhipuAI


def save_embeddings(text_list, save_path, batch_size=64):
    """
    保存文本的向量到本地文件，分批请求最多 64 个文本。
    """
    all_embeddings = []
    
    # 按批次处理文本
    for i in range(0, len(text_list), batch_size):
        chunk = text_list[i:i+batch_size]
        
        # 批量请求文本的嵌入向量
        response = client.embeddings.create(
            model="embedding-3",
            input=chunk
        )
        
        # 提取每个文本的 embedding
        embeddings = [item.embedding for item in response.data]
        all_embeddings.extend(embeddings)  # 累加所有的 embeddings
    
    # 保存到文件
    with open(save_path, 'w', encoding='utf-8') as f:
        json.dump(all_embeddings, f)
    
    print(f"Embeddings saved to {save_path}")

def calculate_similarity(text_list, local_vectors=False):
    """
    批量计算相似度：
      - 输入一个文本列表，其中第一个文本作为基准，其余文本与基准比较
      - 当候选文本超过 64 条时，会分批请求，最后返回所有候选文本与基准文本的相似度（保留4位小数）
    """

    base_text = text_list[0]
    # 先单独请求基准文本的 embedding
    base_response = client.embeddings.create(
        model="embedding-3",
        input=[base_text]
    )
    base_embedding = base_response.data[0].embedding
    base_embedding = np.array(base_embedding)

    all_similarities = []
    if local_vectors:
        if DEBUG_VER == 3:
            print(f'------>local_vectors:{local_vectors}')
        # 使用本地保存的向量
        with open(local_vectors, 'r', encoding='utf-8') as f:
            local_embeddings = json.load(f)
        
        # 确保读取到的向量与候选文本对应
        candidate_embeddings = np.array(local_embeddings)
    else:
        # 批量处理候选文本，每次最多请求64条
        candidate_texts = text_list[1:]
        candidate_embeddings = []
        chunk_size = 64  # 每次最多请求64条
        for i in range(0, len(candidate_texts), chunk_size):
            chunk = candidate_texts[i:i+chunk_size]
            response = client.embeddings.create(
                model="embedding-3",
                input=chunk
            )
            # 提取候选文本的 embedding 并转换为 NumPy 数组
            embeddings = [item.embedding for item in response.data]
            candidate_embeddings.extend(embeddings)
        
        candidate_embeddings = np.array(candidate_embeddings)

    # 计算余弦相似度： dot / (||base|| * ||candidate||)
    dot_products = candidate_embeddings.dot(base_embedding)
    norm_base = np.linalg.norm(base_embedding)
    norm_candidates = np.linalg.norm(candidate_embeddings, axis=1)
    similarities = dot_products / (norm_base * norm_candidates)

    # 保留4位小数，并加入结果列表
    sims = [round(float(sim), 4) for sim in similarities]
    all_similarities.extend(sims)

    return all_similarities