"""
Convert a persisted LlamaIndex PropertyGraphIndex into a PyG Data checkpoint.

Run from the repository root:
    python -m src.graph.pyg_converter
"""

import torch
from torch_geometric.data import Data
from sentence_transformers import SentenceTransformer
from llama_index.core import StorageContext, load_index_from_storage

from src.utils import UTF8LocalFileSystem


def convert_to_pyg(
    persist_dir: str = "./storage_graph",
    output_path: str = "./storage_graph/pyg_data.pt",
):
    print(f"Loading graph from {persist_dir}…")
    try:
        storage_context = StorageContext.from_defaults(
            persist_dir=persist_dir, fs=UTF8LocalFileSystem()
        )
        index = load_index_from_storage(storage_context)
        pg_store = index.property_graph_store
    except Exception as e:
        print(f"Error loading index: {e}")
        return

    nodes_dict = pg_store.graph.nodes
    rel_dict = (
        pg_store.graph.relations.values()
        if isinstance(pg_store.graph.relations, dict)
        else pg_store.graph.relations
    )

    node_id_to_idx: dict = {}
    node_texts: list[str] = []

    for i, (node_id, node_data) in enumerate(nodes_dict.items()):
        node_id_to_idx[node_id] = i

        text_rep = str(node_id)
        if hasattr(node_data, "properties") and node_data.properties:
            text_rep = str(
                node_data.properties.get("name")
                or node_data.properties.get("id")
                or node_id
            )
        if hasattr(node_data, "text") and node_data.text:
            text_rep += " " + node_data.text

        node_texts.append(text_rep)

    print(f"Extracted {len(node_texts)} nodes.")

    print("Encoding nodes with BAAI/bge-small-en-v1.5…")
    embed_model = SentenceTransformer("BAAI/bge-small-en-v1.5")
    embeddings = embed_model.encode(node_texts, show_progress_bar=True, convert_to_tensor=True)

    src_indices: list[int] = []
    dst_indices: list[int] = []

    for rel in rel_dict:
        if hasattr(rel, "source_id") and hasattr(rel, "target_id"):
            src, dst = rel.source_id, rel.target_id
        else:
            src = rel.get("source_id")
            dst = rel.get("target_id")

        if src in node_id_to_idx and dst in node_id_to_idx:
            src_indices.append(node_id_to_idx[src])
            dst_indices.append(node_id_to_idx[dst])

    print(f"Extracted {len(src_indices)} edges.")

    edge_index = torch.tensor([src_indices, dst_indices], dtype=torch.long)
    data = Data(x=embeddings.cpu(), edge_index=edge_index)

    idx_to_node_id = {v: k for k, v in node_id_to_idx.items()}

    print(f"Saving PyG checkpoint to {output_path}…")
    torch.save(
        {
            "pyg_data": data,
            "node_id_to_idx": node_id_to_idx,
            "idx_to_node_id": idx_to_node_id,
        },
        output_path,
    )
    print("Done.")


if __name__ == "__main__":
    convert_to_pyg()