from __future__ import annotations

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


JAKUB_PATH = (
    "pages/students/2016/"
    "jakub_maruniak/README.md"
)


def make_document(
    *,
    path: str = JAKUB_PATH,
    title: str = "Jakub Maruniak",
    categories: list[str] | None = None,
) -> dict:
    return {
        "path": path,
        "title": title,
        "categories": (
            categories
            if categories is not None
            else ["dp2021", "bp2019"]
        ),
        "tags": [
            "spacy",
            "ner",
        ],
        "author": "Daniel Hladek",
        "published": True,
    }


def make_chunk(
    index: int,
    headings: list[str],
    text: str,
    *,
    path: str = JAKUB_PATH,
) -> dict:
    return {
        "chunk_id": (
            f"{path}::chunk-{index}"
        ),
        "document_path": path,
        "title": "Jakub Maruniak",
        "categories": [
            "dp2021",
            "bp2019",
        ],
        "tags": [
            "spacy",
            "ner",
        ],
        "author": "Daniel Hladek",
        "published": True,
        "chunk_index": index,
        "heading_paths": [
            headings
        ],
        "text": text,
        "text_length": len(text),
        "token_count": 1,
        "content_hash": f"hash-{index}",
    }


def test_explicit_dp_proposal_is_found() -> None:
    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Dokument: Jakub Maruniak\n\n"
                "*Návrh na názov DP*:\n\n"
                "Anotácia a rozpoznávanie "
                "pomenovaných entít v "
                "slovenskom jazyku."
            ),
        )
    ]

    title, source_chunk_ids = (
        extract_explicit_proposal_title(
            "diploma_thesis",
            chunks,
        )
    )

    assert title == (
        "Anotácia a rozpoznávanie "
        "pomenovaných entít v "
        "slovenskom jazyku."
    )

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


def test_extract_works_uses_dp_proposal_fallback() -> None:
    document = make_document()

    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Dokument: Jakub Maruniak\n\n"
                "*Rok začiatku štúdia*: 2016\n\n"
                "*Návrh na názov DP*:\n\n"
                "Anotácia a rozpoznávanie "
                "pomenovaných entít v "
                "slovenskom jazyku."
            ),
        ),
        make_chunk(
            1,
            [
                "Jakub Maruniak",
                "Diplomová práca 2021",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomová práca 2021\n\n"
                "Stretnutie 12.3.\n\n"
                "Stav:\n\n"
                "- Anotovanie dát."
            ),
        ),
    ]

    works = extract_works(
        document,
        chunks,
    )

    diploma_works = [
        work
        for work in works
        if work["work_type"]
        == "diploma_thesis"
    ]

    assert len(diploma_works) == 1

    work = diploma_works[0]

    assert work["year"] == 2021
    assert work["title"] == (
        "Anotácia a rozpoznávanie "
        "pomenovaných entít v "
        "slovenskom jazyku."
    )

    assert work["source_chunk_ids"] == [
        f"{JAKUB_PATH}::chunk-0",
        f"{JAKUB_PATH}::chunk-1",
    ]


def test_existing_work_title_has_priority() -> None:
    document = make_document()

    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Dokument: Jakub Maruniak\n\n"
                "Návrh na názov DP: "
                "Starší pracovný názov"
            ),
        ),
        make_chunk(
            1,
            [
                "Jakub Maruniak",
                "Diplomová práca 2021",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomová práca 2021\n\n"
                "Názov: Finálny názov práce"
            ),
        ),
    ]

    works = extract_works(
        document,
        chunks,
    )

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

    assert work["title"] == (
        "Finálny názov práce"
    )

    assert work["source_chunk_ids"] == [
        f"{JAKUB_PATH}::chunk-1"
    ]


def test_ambiguous_dp_proposals_are_not_used() -> None:
    document = make_document()

    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Návrh na názov DP: "
                "Prvý návrh"
            ),
        ),
        make_chunk(
            1,
            ["Jakub Maruniak"],
            (
                "Návrh na názov DP: "
                "Druhý návrh"
            ),
        ),
        make_chunk(
            2,
            [
                "Jakub Maruniak",
                "Diplomová práca 2021",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomová práca 2021\n\n"
                "Stretnutie 12.3."
            ),
        ),
    ]

    works = extract_works(
        document,
        chunks,
    )

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

    assert work["title"] is None


def test_proposal_is_not_used_for_diploma_project() -> None:
    document = make_document()

    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Návrh na názov DP: "
                "Názov diplomovej práce"
            ),
        ),
        make_chunk(
            1,
            [
                "Jakub Maruniak",
                "Diplomový projekt 2020",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomový projekt 2020\n\n"
                "Stretnutie 9.3.2020\n\n"
                "Úlohy:\n\n"
                "- Pripraviť dáta."
            ),
        ),
    ]

    works = extract_works(
        document,
        chunks,
    )

    work = next(
        work
        for work in works
        if work["work_type"]
        == "diploma_project"
    )

    assert work["title"] is None


def test_proposal_is_not_shared_between_multiple_dp_works() -> None:
    document = make_document(
        categories=[
            "dp2021",
            "dp2022",
        ]
    )

    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Návrh na názov DP: "
                "Jeden návrh"
            ),
        ),
        make_chunk(
            1,
            [
                "Jakub Maruniak",
                "Diplomová práca 2021",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomová práca 2021\n\n"
                "Stretnutie 1.1.2021"
            ),
        ),
        make_chunk(
            2,
            [
                "Jakub Maruniak",
                "Diplomová práca 2022",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomová práca 2022\n\n"
                "Stretnutie 1.1.2022"
            ),
        ),
    ]

    works = extract_works(
        document,
        chunks,
    )

    diploma_works = [
        work
        for work in works
        if work["work_type"]
        == "diploma_thesis"
    ]

    assert len(diploma_works) == 2
    assert all(
        work["title"] is None
        for work in diploma_works
    )


def test_graph_payload_contains_fallback_title() -> None:
    document = make_document()

    chunks = [
        make_chunk(
            0,
            ["Jakub Maruniak"],
            (
                "Dokument: Jakub Maruniak\n\n"
                "Rok začiatku štúdia: 2016\n\n"
                "Návrh na názov DP:\n\n"
                "Anotácia a rozpoznávanie "
                "pomenovaných entít v "
                "slovenskom jazyku."
            ),
        ),
        make_chunk(
            1,
            [
                "Jakub Maruniak",
                "Diplomová práca 2021",
            ],
            (
                "Dokument: Jakub Maruniak\n"
                "Sekcia: Diplomová práca 2021\n\n"
                "Stretnutie 12.3."
            ),
        ),
    ]

    payload = build_graph_payload(
        [document],
        chunks,
    )

    diploma_work = next(
        work
        for work in payload["works"]
        if work["work_type"]
        == "diploma_thesis"
    )

    assert diploma_work["title"] == (
        "Anotácia a rozpoznávanie "
        "pomenovaných entít v "
        "slovenskom jazyku."
    )

    assert diploma_work[
        "source_chunk_ids"
    ] == [
        f"{JAKUB_PATH}::chunk-0",
        f"{JAKUB_PATH}::chunk-1",
    ]