"""
Build a LlamaIndex PropertyGraphIndex from PubMedQA and persist it to disk.

Run from the repository root:
    python -m src.indexing.kg_builder
"""

from typing import List
from datasets import load_dataset
from llama_index.core import Document, PropertyGraphIndex, Settings
from llama_index.core.indices.property_graph import SimpleLLMPathExtractor
from llama_index.llms.llama_cpp import LlamaCPP

from src.utils import MODEL_REPO_ID, MODEL_FILENAME, UTF8LocalFileSystem, resolve_gguf_model_path

MAX_DOCS = 10


def split_text(text: str, chunk_size: int = 150):
    words = text.split()
    for i in range(0, len(words), chunk_size):
        yield " ".join(words[i : i + chunk_size])


def build_kg(persist_dir: str = "./storage_graph"):
    print("Loading PubMedQA dataset…")
    dataset = load_dataset("qiaojin/PubMedQA", "pqa_labeled", split="train")

    documents: List[Document] = []
    count = 0

    for item in dataset:
        if count >= MAX_DOCS:
            break

        contexts = item["context"]["contexts"]
        labels = item["context"]["labels"]
        meshes = item["context"]["meshes"]

        for ctx, label in zip(contexts, labels):
            if not ctx or not ctx.strip():
                continue
            for chunk in split_text(ctx):
                documents.append(
                    Document(
                        text=chunk.strip(),
                        metadata={
                            "pubid": item["pubid"],
                            "section": label,
                            "mesh_terms": ", ".join(meshes) if meshes else "",
                        },
                    )
                )
        count += 1

    print(f"Total chunks prepared: {len(documents)}")

    print("Initialising LLM for KG extraction…")
    llm = LlamaCPP(
        model_path=resolve_gguf_model_path(),
        temperature=0.0,
        max_new_tokens=256,
        context_window=2048,
        model_kwargs={"n_threads": 4, "n_ctx": 2048, "n_batch": 32},
        verbose=False,
    )
    Settings.llm = llm

    kg_extractor = SimpleLLMPathExtractor(llm=llm, max_paths_per_chunk=2)

    print("Building PropertyGraphIndex (this may take several minutes)…")
    index = PropertyGraphIndex.from_documents(
        documents,
        kg_extractors=[kg_extractor],
        show_progress=True,
    )

    print(f"Persisting graph to {persist_dir}…")
    index.storage_context.persist(persist_dir, fs=UTF8LocalFileSystem())
    print("Done.")


if __name__ == "__main__":
    build_kg("./storage_graph")