"""
Pre-compute KB section embeddings + build FAISS index.
Run during Docker build (and locally). ONNX export optional (--export-onnx).
"""
import os, sys, json, time, argparse
sys.stdout.reconfigure(encoding="utf-8")

parser = argparse.ArgumentParser()
parser.add_argument("--export-onnx", action="store_true")
args = parser.parse_args()

import numpy as np
import faiss
import torch
from sentence_transformers import SentenceTransformer

os.makedirs("embeddings", exist_ok=True)

# 1. Load model
t0 = time.time()
model = SentenceTransformer("BAAI/bge-m3", trust_remote_code=True)
print(f"Model loaded in {time.time()-t0:.1f}s")

# 2. Load KB sections
paths = ["knowledge_base.json", "knowledge_base_insoles.json"]
sections = []
for path in paths:
    with open(path, encoding="utf-8") as f:
        data = json.load(f)
    for name, sec in data["sections"].items():
        sections.append({
            "id": f"{path}:{name}",
            "title": sec["title"],
            "content": sec["content"],
        })
print(f"Loaded {len(sections)} sections")

# 3. Embed all sections
texts = [f"{s['title']}\n\n{s['content']}" for s in sections]
t1 = time.time()
embeddings = model.encode(texts, normalize_embeddings=True, show_progress_bar=True)
print(f"Embedded {len(embeddings)} sections in {time.time()-t1:.1f}s")

# 4. Build and save FAISS index
index = faiss.IndexFlatIP(embeddings.shape[1])
index.add(embeddings.astype(np.float32))
faiss.write_index(index, "embeddings/index.faiss")
with open("embeddings/sections.json", "w", encoding="utf-8") as f:
    json.dump(sections, f, ensure_ascii=False, indent=2)
np.save("embeddings/embeddings.npy", embeddings)
print(f"FAISS: {index.ntotal} vectors x {embeddings.shape[1]} dim")
print(f"  index.faiss: {os.path.getsize('embeddings/index.faiss')/1024:.1f} KB")
print(f"  sections.json: {os.path.getsize('embeddings/sections.json')/1024:.1f} KB")

if not args.export_onnx:
    print("Skipping ONNX export (not needed at runtime)")
    print(f"Done in {time.time()-t0:.1f}s")
    sys.exit(0)

# 5. Export to ONNX via torch
t2 = time.time()
print("Exporting transformer to ONNX...")
raw_model = model[0].auto_model.eval()
tok = model.tokenizer
dummy = tok("test query", return_tensors="pt", padding=True, truncation=True, max_length=256)
with torch.no_grad():
    torch.onnx.export(
        raw_model,
        (dummy["input_ids"], dummy["attention_mask"]),
        "embeddings/bge-m3.onnx",
        input_names=["input_ids", "attention_mask"],
        output_names=["last_hidden_state"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "seq"},
            "attention_mask": {0: "batch", 1: "seq"},
        },
        opset_version=14,
    )
onnx_size_mb = os.path.getsize("embeddings/bge-m3.onnx") / 1024 / 1024
print(f"ONNX exported in {time.time()-t2:.1f}s, {onnx_size_mb:.1f} MB")
print(f"Done in {time.time()-t0:.1f}s")