import json
from pathlib import Path
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from data_io import load_clauses, write_jsonl

def build_one(model_name, out_dir, ru, kz, batch_size=64):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    ru_texts = [x["text"] for x in ru]
    kz_texts = [x["text"] for x in kz]

    ru_meta = [{"pos": i, "id": ru[i]["id"], "lang": "ru", "meta": ru[i]["meta"], "text": ru[i]["text"]} for i in range(len(ru))]
    kz_meta = [{"pos": i, "id": kz[i]["id"], "lang": "kz", "meta": kz[i]["meta"], "text": kz[i]["text"]} for i in range(len(kz))]

    model = SentenceTransformer(model_name)

    ru_emb = model.encode(ru_texts, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True).astype(np.float32)
    kz_emb = model.encode(kz_texts, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True).astype(np.float32)

    ru_index = faiss.IndexFlatIP(ru_emb.shape[1])
    ru_index.add(ru_emb)
    faiss.write_index(ru_index, str(out_dir / "ru.faiss"))

    kz_index = faiss.IndexFlatIP(kz_emb.shape[1])
    kz_index.add(kz_emb)
    faiss.write_index(kz_index, str(out_dir / "kz.faiss"))

    write_jsonl(str(out_dir / "ru_meta.jsonl"), ru_meta)
    write_jsonl(str(out_dir / "kz_meta.jsonl"), kz_meta)

    info = {
        "model": model_name,
        "ru_count": len(ru_texts),
        "kz_count": len(kz_texts),
        "ru_index": str(out_dir / "ru.faiss"),
        "kz_index": str(out_dir / "kz.faiss"),
        "ru_meta": str(out_dir / "ru_meta.jsonl"),
        "kz_meta": str(out_dir / "kz_meta.jsonl"),
    }
    (out_dir / "index_info.json").write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")

def main():
    clauses_path = "data/clauses_constitution_ru_kz.jsonl"
    ru, kz = load_clauses(clauses_path)

    base_root = Path("artifacts/indexes")
    base_root.mkdir(parents=True, exist_ok=True)

    models = [
        ("mpnet_base", "paraphrase-multilingual-mpnet-base-v2"),
        ("labse", "sentence-transformers/LaBSE"),
    ]

    finetuned_path = Path("artifacts/models/finetuned_mpnet")
    if finetuned_path.exists():
        models.append(("finetuned", str(finetuned_path)))

    built = []
    for alias, model_name in models:
        out_dir = base_root / alias
        build_one(model_name, out_dir, ru, kz, batch_size=64)
        built.append({"alias": alias, "dir": str(out_dir), "model": model_name})

    (base_root / "built_indexes.json").write_text(json.dumps(built, ensure_ascii=False, indent=2), encoding="utf-8")

if __name__ == "__main__":
    main()