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"
    ] == []