"""Rebuild LightRAG vector stores from the knowledge graph and text chunks.

The vector files (vdb_entities.json, vdb_relationships.json, vdb_chunks.json)
are derived data: every record is an embedding of an entity or relation stored
in graph_chunk_entity_relation.graphml, or of a chunk in
kv_store_text_chunks.json. If they get corrupted (e.g. the process
ran out of memory while saving), they can be recomputed locally without the
LLM. Records use the same ids, fields and content format as LightRAG
(lightrag/operate.py), so the server loads them as if it wrote them itself.

Stop the LightRAG server before running this script.

    python scripts/lightrag_ingest/rebuild_vdb.py --storage lightrag/rag_storage_qwen36
"""

from __future__ import annotations

import os

# Prevents the PyTorch CPU allocator cache from growing with varying input shapes.
os.environ.setdefault("LRU_CACHE_CAPACITY", "1")

import argparse
import base64
import json
import shutil
import sys
import time
import zlib
from pathlib import Path

import networkx as nx
import numpy as np
from nano_vectordb import NanoVectorDB
from sentence_transformers import SentenceTransformer

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "lightrag"))
from lightrag.utils import compute_mdhash_id  # noqa: E402

MODEL_NAME = "paraphrase-multilingual-MiniLM-L12-v2"  # same model as embedding_server.py

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")


def entity_records(graph: nx.Graph) -> dict[str, dict]:
    records = {}
    for name, data in graph.nodes(data=True):
        records[compute_mdhash_id(name, prefix="ent-")] = {
            "entity_name": name,
            "content": f"{name}\n{data.get('description', '')}",
            "source_id": data.get("source_id", ""),
            "file_path": data.get("file_path", ""),
        }
    return records


def relation_records(graph: nx.Graph) -> dict[str, dict]:
    records = {}
    for u, v, data in graph.edges(data=True):
        src, tgt = (u, v) if u <= v else (v, u)
        keywords = data.get("keywords", "")
        records[compute_mdhash_id(src + tgt, prefix="rel-")] = {
            "src_id": src,
            "tgt_id": tgt,
            "content": f"{keywords}\t{src}\n{tgt}\n{data.get('description', '')}",
            "source_id": data.get("source_id", ""),
            "file_path": data.get("file_path", ""),
        }
    return records


def chunk_records(storage: Path) -> dict[str, dict]:
    chunks = json.loads((storage / "kv_store_text_chunks.json").read_text(encoding="utf-8"))
    return {
        chunk_id: {
            "full_doc_id": chunk.get("full_doc_id", ""),
            "content": chunk.get("content", ""),
            "file_path": chunk.get("file_path", ""),
        }
        for chunk_id, chunk in chunks.items()
    }


def build_store(path: Path, records: dict[str, dict], model: SentenceTransformer, batch_size: int) -> None:
    ids = list(records)
    contents = [records[i]["content"] for i in ids]
    print(f"Embedding {len(ids)} records for {path.name} ...")
    started = time.time()
    vectors = model.encode(
        contents,
        batch_size=batch_size,
        normalize_embeddings=True,
        show_progress_bar=True,
        convert_to_numpy=True,
    ).astype(np.float32)
    print(f"  done in {time.time() - started:.0f}s")

    now = int(time.time())
    datas = []
    for record_id, vector in zip(ids, vectors):
        encoded = base64.b64encode(zlib.compress(vector.astype(np.float16).tobytes())).decode("utf-8")
        datas.append({
            "__id__": record_id,
            "__created_at__": now,
            **records[record_id],
            "vector": encoded,
            "__vector__": vector,
        })

    tmp = path.with_suffix(".rebuilt.json")
    if tmp.exists():
        tmp.unlink()
    db = NanoVectorDB(vectors.shape[1], storage_file=str(tmp))
    db.upsert(datas=datas)
    db.save()

    if path.exists():
        backup = path.with_suffix(".broken.json")
        shutil.move(str(path), str(backup))
        print(f"  old file moved to {backup.name}")
    shutil.move(str(tmp), str(path))
    print(f"  saved {path}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--storage", type=Path, default=ROOT / "lightrag" / "rag_storage_qwen36")
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument(
        "--only", choices=("entities", "relationships", "chunks"), help="Rebuild only one store."
    )
    args = parser.parse_args()

    model = SentenceTransformer(MODEL_NAME)
    if args.only in (None, "entities", "relationships"):
        graph_file = args.storage / "graph_chunk_entity_relation.graphml"
        print(f"Loading graph {graph_file} ...")
        graph = nx.read_graphml(graph_file)
        print(f"  {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges")
        if args.only in (None, "entities"):
            build_store(args.storage / "vdb_entities.json", entity_records(graph), model, args.batch_size)
        if args.only in (None, "relationships"):
            build_store(args.storage / "vdb_relationships.json", relation_records(graph), model, args.batch_size)
    if args.only in (None, "chunks"):
        build_store(args.storage / "vdb_chunks.json", chunk_records(args.storage), model, args.batch_size)
    print("Finished.")


if __name__ == "__main__":
    main()