"""
Step 1: 加载中文 embedding 模型
Step 2: 为 _chunks.jsonl 全部 22,781 条计算向量
Step 3: 写入 FAISS 索引（含分类过滤元数据）

输出：
  _vectors.np        — 22,781 × dim 向量矩阵
  _vector_meta.json  — 每行 chunk 的 metadata（含 category）
  _faiss_index.bin   — FAISS 索引文件
"""
import json, os, sys
import numpy as np
from pathlib import Path
from collections import defaultdict

BASE = Path(r"C:\code\cfc\docs\菌群知识库")
CHUNKS_FILE = BASE / "_chunks.jsonl"
VEC_NP = BASE / "_vectors.npy"
VEC_META = BASE / "_vector_meta.json"
FAISS_INDEX = BASE / "_faiss_index.bin"

EMBED_MODEL = "BAAI/bge-small-zh-v1.5"
MAX_SEQ_LEN = 256
BATCH_SIZE = 128

print("=== Step 1: 加载中文 embedding 模型 ===")
from sentence_transformers import SentenceTransformer
print(f"  加载模型: {EMBED_MODEL} ...")
model = SentenceTransformer(EMBED_MODEL)
vec_dim = model.get_sentence_embedding_dimension()
print(f"  向量维度: {vec_dim}")
print("  模型加载完成")

print("\n=== Step 2: 读取 JSONL 并计算向量 ===")
chunks_data = []
meta_list = []

with open(CHUNKS_FILE, "r", encoding="utf-8") as f:
    for line in f:
        d = json.loads(line)
        chunks_data.append(d)

print(f"  共 {len(chunks_data)} 条 chunk")

# 分批计算向量
print(f"  开始计算向量（batch_size={BATCH_SIZE}）...")
all_vectors = []
for i in range(0, len(chunks_data), BATCH_SIZE):
    batch = chunks_data[i:i+BATCH_SIZE]
    texts = [c["content"] for c in batch]
    vecs = model.encode(texts, batch_size=BATCH_SIZE, show_progress_bar=(i % 2000 == 0), normalize_embeddings=True)
    all_vectors.append(vecs)
    if (i // BATCH_SIZE + 1) % 10 == 0:
        print(f"    processed {min(i+BATCH_SIZE, len(chunks_data))}/{len(chunks_data)} chunks...")

vectors = np.concatenate(all_vectors, axis=0)
vectors = np.ascontiguousarray(vectors, dtype=np.float32)
print(f"  向量矩阵: {vectors.shape}")

# 保存向量
np.save(VEC_NP, vectors)
print(f"  向量已保存: {VEC_NP}")

# 保存 metadata（顺序与向量一致）
for c in chunks_data:
    meta_list.append({
        "id": c["id"],
        "title": c["metadata"]["title"],
        "postId": c["metadata"]["postId"],
        "category": c["metadata"]["category"],
        "source_url": c["metadata"].get("source_url", ""),
        "topics": c["metadata"].get("topics", []),
        "chunk_index": c["metadata"]["chunk_index"],
        "chunk_total": c["metadata"]["chunk_total"],
        "filename": c["metadata"]["filename"],
        "content_preview": c["content"][:120],
    })

with open(VEC_META, "w", encoding="utf-8") as f:
    json.dump(meta_list, f, ensure_ascii=False)
print(f"  元数据已保存: {VEC_META}")

# 分类过滤索引（category → chunk_indices）
cat_indices = defaultdict(list)
for idx, m in enumerate(meta_list):
    cat_indices[m["category"]].append(idx)

cat_filter = json.dumps({k: v for k, v in cat_indices.items()}, ensure_ascii=False)
with open(BASE / "_cat_filter.json", "w", encoding="utf-8") as f:
    json.dump({k: v for k, v in cat_indices.items()}, f, ensure_ascii=False)

print("\n=== Step 3: 写入 FAISS 索引 ===")
import faiss

# 创建 HNSW 索引（带分类过滤支持）
print(f"  创建 FAISS IDMap + HNSW 索引...")
idx_flat = faiss.IndexFlatIP(vec_dim)  # 内积（已归一化向量）
idx = faiss.IndexIDMap2(idx_flat)

# 添加向量（IDs = 0..N-1）
idx.add_with_ids(vectors, np.arange(len(vectors)))
print(f"  索引已添加 {idx.ntotal} 条向量")

# 保存 FAISS 索引
faiss.write_index(idx, str(FAISS_INDEX))
print(f"  FAISS 索引已保存: {FAISS_INDEX}")

# 统计分类覆盖
print("\n=== 分类统计 ===")
for cat, indices in sorted(cat_indices.items()):
    print(f"  {cat}: {len(indices)} chunks")

print(f"\n✅ 完成！{len(chunks_data)} chunks → {vectors.shape[1]}d 向量 → FAISS 索引")