"""
构建菜品名向量索引
将labels.py中的菜品名编码为向量并构建FAISS索引（一次性任务）
"""
import argparse
import json
import os
import sys
import time
from typing import List

import faiss
import numpy as np

# 兼容：支持直接运行脚本或用 -m 模块方式运行
if __name__ == "__main__" and __package__ is None:
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from exp_multimodal.labels import build_labels
from exp_multimodal.text_embedder import OllamaEmbedder


def build_index(
    dish_names: List[str],
    embedder: OllamaEmbedder,
    output_dir: str,
    batch_size: int = 100,
) -> None:
    """
    构建FAISS索引
    
    参数：
        dish_names: 菜品名列表
        embedder: OllamaEmbedder实例
        output_dir: 输出目录
        batch_size: 批量编码大小
    """
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"[BuildIndex] Total dishes={len(dish_names)} batch_size={batch_size}")
    
    # 批量编码（带重试机制）
    all_embeddings = []
    for i in range(0, len(dish_names), batch_size):
        batch = dish_names[i:i+batch_size]
        batch_num = i//batch_size + 1
        total_batches = (len(dish_names)-1)//batch_size + 1
        print(f"[BuildIndex] Encoding batch {batch_num}/{total_batches} (size={len(batch)})...")
        
        # 重试机制：最多3次，指数退避
        max_retries = 3
        for retry in range(max_retries):
            try:
                batch_embs = embedder.encode(batch)
                all_embeddings.append(batch_embs)
                break  # 成功则跳出重试循环
            except Exception as e:
                if retry < max_retries - 1:
                    wait_time = 5 * (retry + 1)  # 5s, 10s, 15s
                    print(f"[BuildIndex] Batch {batch_num} failed (attempt {retry+1}/{max_retries}): {e}")
                    print(f"[BuildIndex] Retrying in {wait_time} seconds...")
                    time.sleep(wait_time)
                else:
                    print(f"[BuildIndex] Batch {batch_num} failed after {max_retries} attempts: {e}")
                    raise
    
    # 合并所有向量
    embeddings = np.vstack(all_embeddings)
    print(f"[BuildIndex] Concatenated embeddings shape={embeddings.shape}")
    
    # 归一化向量（用于余弦相似度）
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    embeddings = embeddings / (norms + 1e-8)
    print(f"[BuildIndex] Normalized embeddings")
    
    # 构建FAISS索引（IndexFlatIP = 内积索引，适合归一化后的向量）
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings.astype(np.float32))
    print(f"[BuildIndex] Built FAISS index dim={dim} ntotal={index.ntotal}")
    
    # 保存文件
    names_path = os.path.join(output_dir, "dish_names.json")
    embeddings_path = os.path.join(output_dir, "dish_embeddings.npy")
    index_path = os.path.join(output_dir, "faiss_index.bin")
    
    with open(names_path, "w", encoding="utf-8") as f:
        json.dump(dish_names, f, ensure_ascii=False, indent=2)
    print(f"[BuildIndex] Saved dish names to {names_path}")
    
    np.save(embeddings_path, embeddings)
    print(f"[BuildIndex] Saved embeddings to {embeddings_path}")
    
    faiss.write_index(index, index_path)
    print(f"[BuildIndex] Saved FAISS index to {index_path}")
    
    print(f"[BuildIndex] ✅ Index build complete! Output dir: {output_dir}")


def main():
    ap = argparse.ArgumentParser(description="构建菜品名向量索引")
    ap.add_argument(
        "--mode",
        choices=["dish", "whole", "processed"],
        default="dish",
        help="数据集模式（默认: dish）"
    )
    ap.add_argument(
        "--output",
        default=None,
        help="输出目录（默认: faiss_vector_db/<mode>_names）"
    )
    ap.add_argument(
        "--embedder_url",
        default="http://192.168.1.250:11434",
        help="Ollama服务地址"
    )
    ap.add_argument(
        "--embedder_model",
        default="quentinz/bge-large-zh-v1.5",
        help="Embedding模型名称"
    )
    ap.add_argument(
        "--batch_size",
        type=int,
        default=100,
        help="批量编码大小"
    )
    
    args = ap.parse_args()
    
    # 确定输出目录
    if args.output:
        output_dir = args.output
    else:
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        output_dir = os.path.join(project_root, "faiss_vector_db", f"{args.mode}_names")
    
    print(f"[Main] mode={args.mode} output_dir={output_dir}")
    print(f"[Main] embedder_url={args.embedder_url} model={args.embedder_model}")
    
    # 构建菜品名列表
    print(f"[Main] Building labels from mode={args.mode}...")
    dish_names = build_labels(args.mode, alias_map_path=None)
    print(f"[Main] Built {len(dish_names)} dish names")
    
    if not dish_names:
        print("[Main] ❌ No dish names found, abort")
        return
    
    # 初始化Embedder
    print(f"[Main] Initializing OllamaEmbedder...")
    embedder = OllamaEmbedder(
        base_url=args.embedder_url,
        model=args.embedder_model,
        timeout=180,  # 3分钟超时
    )
    
    # 构建索引
    print(f"[Main] Building FAISS index...")
    build_index(
        dish_names=dish_names,
        embedder=embedder,
        output_dir=output_dir,
        batch_size=args.batch_size,
    )
    
    print(f"[Main] ✅ All done!")


if __name__ == "__main__":
    main()