from __future__ import annotations

from tools.graphrag.graphrag import (
    build_graph_payload,
    extract_works,
)

DARIUS_PATH = "pages/students/2016/" "darius_lindvai/README.md"


def make_darius_document() -> dict:
    return {
        "path": DARIUS_PATH,
        "title": "Dárius Lindvai",
        "categories": [
            "dp2021",
            "bp2019",
        ],
        "tags": [
            "nn",
            "interpunction",
            "nlp",
        ],
        "author": "Daniel Hladek",
        "published": True,
    }


def make_root_chunk() -> dict:
    text = (
        "Dokument: Dárius Lindvai\n\n"
        "Rok začiatku štúdia: 2016\n\n"
        "Repozitár so zdrojovými kódmi\n\n"
        "Názov: Obnovenie interpunkcie pomocou "
        "hlbokých neurónových sietí\n\n"
        "1. Vypracujte prehľad metód.\n"
        "2. Vyberte vhodnú metódu."
    )

    return {
        "chunk_id": (f"{DARIUS_PATH}::chunk-0"),
        "document_path": DARIUS_PATH,
        "title": "Dárius Lindvai",
        "categories": [
            "dp2021",
            "bp2019",
        ],
        "tags": [
            "nn",
            "interpunction",
            "nlp",
        ],
        "author": "Daniel Hladek",
        "published": True,
        "chunk_index": 0,
        "heading_paths": [
            [
                "Dárius Lindvai",
            ]
        ],
        "text": text,
        "text_length": len(text),
        "token_count": 1,
        "content_hash": "darius-root",
    }


def test_root_title_creates_missing_latest_thesis() -> None:
    works = extract_works(
        make_darius_document(),
        [
            make_root_chunk(),
        ],
    )

    diploma = [work for work in works if (work["work_type"] == "diploma_thesis")]

    bachelor = [work for work in works if (work["work_type"] == "bachelor_thesis")]

    assert len(diploma) == 1

    assert bachelor == []

    thesis = diploma[0]

    assert thesis["year"] == 2021

    assert thesis["heading"] == "Diplomová práca 2021"

    assert thesis["title"] == ("Obnovenie interpunkcie pomocou " "hlbokých neurónových sietí")

    assert thesis["source_chunk_ids"] == [f"{DARIUS_PATH}::chunk-0"]


def test_graph_payload_contains_darius_thesis() -> None:
    payload = build_graph_payload(
        [
            make_darius_document(),
        ],
        [
            make_root_chunk(),
        ],
    )

    assert len(payload["theses"]) == 1

    thesis = payload["theses"][0]

    assert thesis["thesis_type"] == "diploma_thesis"

    assert thesis["year"] == 2021

    assert thesis["title"] == ("Obnovenie interpunkcie pomocou " "hlbokých neurónových sietí")

    assert thesis["source_document"] == DARIUS_PATH

    assert payload["works"] == []