from __future__ import annotations

import sqlite3
from pathlib import Path

from scripts.rag_supplemental_expansion import (
    canonical_student_main_path,
    expand_results_with_supplemental_sources,
    is_multi_document_query,
)


def create_db(tmp_path: Path) -> Path:
    db = tmp_path / "supplemental.sqlite"

    with sqlite3.connect(db) as conn:
        conn.execute(
            """
            CREATE TABLE documents (
                id INTEGER PRIMARY KEY,
                path TEXT NOT NULL,
                title TEXT,
                author TEXT,
                published INTEGER
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE chunks (
                id INTEGER PRIMARY KEY,
                chunk_id TEXT NOT NULL,
                document_path TEXT NOT NULL,
                title TEXT,
                author TEXT,
                published INTEGER,
                chunk_index INTEGER,
                heading_paths_json TEXT,
                text TEXT
            )
            """
        )

    return db


def add_document(
    db: Path,
    *,
    path: str,
    title: str,
    text: str,
    chunk_index: int = 0,
    published: bool = True,
) -> None:
    with sqlite3.connect(db) as conn:
        conn.execute(
            """
            INSERT OR IGNORE INTO documents(
                path,
                title,
                author,
                published
            ) VALUES (?, ?, ?, ?)
            """,
            (
                path,
                title,
                "Daniel Hladek",
                1 if published else 0,
            ),
        )
        conn.execute(
            """
            INSERT INTO chunks(
                chunk_id,
                document_path,
                title,
                author,
                published,
                chunk_index,
                heading_paths_json,
                text
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                f"{path}::chunk-{chunk_index}",
                path,
                title,
                "Daniel Hladek",
                1 if published else 0,
                chunk_index,
                "[]",
                text,
            ),
        )


def base_result(path: str, title: str) -> dict:
    return {
        "chunk_id": f"{path}::chunk-9",
        "document_path": path,
        "title": title,
        "author": "Daniel Hladek",
        "published": True,
        "chunk_index": 9,
        "heading_paths": [],
        "text": "menej relevantný úsek",
        "source_url": "https://example.invalid/base",
    }


def test_canonical_student_main_path_maps_child_document() -> None:
    assert canonical_student_main_path(
        "pages/students/2016/jakub_maruniak/dp2021/README.md"
    ) == "pages/students/2016/jakub_maruniak/README.md"


def test_exact_person_title_anchors_own_student_document(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    path = "pages/students/2018/ondrej_megela/README.md"
    add_document(
        db,
        path=path,
        title="Ondrej Megela",
        text=(
            "Ondrej Megela. Neurónové siete a neurónové "
            "jazykové modelovanie typu BERT."
        ),
    )
    wrong_path = "pages/students/2016/patrik_pavlisin/README.md"
    add_document(
        db,
        path=wrong_path,
        title="Patrik Pavlišin",
        text="Ondrej Megela je spomenutý pri neurónových sieťach.",
    )

    result = expand_results_with_supplemental_sources(
        db,
        "Ondrej Megela neurónové siete dokument študent",
        [base_result(wrong_path, "Patrik Pavlišin")],
        published_only=True,
        limit=5,
    )

    assert result[0]["document_path"] == path
    assert result[0]["title"] == "Ondrej Megela"


def test_student_family_evidence_includes_child_document(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    main = "pages/students/2016/jakub_maruniak/README.md"
    child = "pages/students/2016/jakub_maruniak/dp2021/README.md"
    add_document(
        db,
        path=main,
        title="Jakub Maruniak",
        text="Jakub Maruniak. Rok začiatku štúdia: 2016.",
    )
    add_document(
        db,
        path=child,
        title="Diplomová práca Jakuba Maruniaka",
        text=(
            "Jakub Maruniak. Diplomová práca 2021. "
            "Názov: Anotácia a rozpoznávanie pomenovaných entít "
            "v slovenskom jazyku."
        ),
    )

    result = expand_results_with_supplemental_sources(
        db,
        "Jakub Maruniak diplomová práca názov",
        [base_result(main, "Jakub Maruniak")],
        published_only=True,
        limit=5,
    )

    blocks = result[0].get("family_evidence_blocks") or []
    assert blocks
    assert any(block["document_path"] == child for block in blocks)
    assert any("Anotácia" in block["focus_text"] for block in blocks)


def test_reverse_lookup_adds_document_with_exact_work_phrase(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    expected = "pages/students/2016/patrik_pavlisin/README.md"
    add_document(
        db,
        path=expected,
        title="Patrik Pavlišin",
        text=(
            "Bakalárska práca 2020. "
            "Názov bakalárskej práce: Metódy automatického prekladu."
        ),
    )
    add_document(
        db,
        path="pages/students/2020/vladyslav_krupko/README.md",
        title="Vladyslav Krupko",
        text="Strojový preklad je súčasťou úloh.",
    )

    result = expand_results_with_supplemental_sources(
        db,
        "Metódy automatického prekladu študent dokument téma",
        [],
        published_only=True,
        limit=5,
    )

    assert result[0]["document_path"] == expected


def test_bilingual_reverse_lookup_finds_internship(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    expected = "pages/interns/bogdan_paul_chis/README.md"
    add_document(
        db,
        path=expected,
        title="Bogdan Paul Chiș",
        text=(
            "Topic: (multilingual) Triplet extraction from medical data. "
            "Goal: Construct a knowledge graph from medical package inserts."
        ),
    )
    add_document(
        db,
        path="pages/topics/otvorene/README.md",
        title="Otvorené projekty",
        text="Multilingual extraction from medical documents is an open topic.",
    )

    result = expand_results_with_supplemental_sources(
        db,
        "multilinguálna extrakcia trojíc z medicínskych dát",
        [],
        published_only=True,
        limit=5,
    )

    assert result[0]["document_path"] == expected


def test_exact_topic_page_title_is_promoted(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    topic = "pages/topics/translation/README.md"
    add_document(
        db,
        path=topic,
        title="Strojový preklad",
        text="Štatistický strojový preklad. Neurónový strojový preklad.",
    )
    add_document(
        db,
        path="pages/students/2016/patrik_pavlisin/README.md",
        title="Patrik Pavlišin",
        text="Metódy automatického a strojového prekladu.",
    )

    result = expand_results_with_supplemental_sources(
        db,
        "Strojový preklad projektová stránka informačná stránka",
        [],
        published_only=True,
        limit=5,
    )

    assert result[0]["document_path"] == topic


def test_generic_question_prefers_root_topic_page(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    root = "pages/topics/question/README.md"
    child = "pages/topics/question/validacie/README.md"
    add_document(
        db,
        path=root,
        title="Question Answering",
        text=(
            "Anotácia otázok. Pre každý odsek má vzniknúť 5 otázok. "
            "Používa sa anotačný nástroj Prodigy."
        ),
    )
    add_document(
        db,
        path=child,
        title="Validácia otázok a odpovedí",
        text="Anotácia otázok a validácia odsekov.",
    )

    result = expand_results_with_supplemental_sources(
        db,
        "Koľko otázok má vzniknúť pre každý odsek pri anotácii otázok?",
        [],
        published_only=True,
        limit=5,
    )

    assert result[0]["document_path"] == root


def test_multi_document_query_adds_multiple_sources(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    paths = [
        "pages/students/2022/jan_ptak/README.md",
        "pages/students/2022/jan_malinovsky/README.md",
        "pages/students/2021/matej_scislak/README.md",
    ]

    for index, path in enumerate(paths):
        add_document(
            db,
            path=path,
            title=f"Student {index}",
            text="RAG a GraphRAG, retrieval augmented generation.",
        )

    result = expand_results_with_supplemental_sources(
        db,
        "Nájdi viacero dokumentov, ktoré riešia RAG alebo GraphRAG.",
        [],
        published_only=True,
        limit=10,
    )

    assert is_multi_document_query(
        "Nájdi viacero dokumentov, ktoré riešia RAG alebo GraphRAG."
    )
    assert len(result) >= 2
    assert len({item["document_path"] for item in result}) >= 2


def test_unpublished_supplement_is_excluded(tmp_path: Path) -> None:
    db = create_db(tmp_path)
    add_document(
        db,
        path="pages/topics/secret/README.md",
        title="Tajný projekt",
        text="Tajný projekt a unikátny výraz xyzabc.",
        published=False,
    )

    result = expand_results_with_supplemental_sources(
        db,
        "Tajný projekt xyzabc",
        [],
        published_only=True,
        limit=5,
    )

    assert result == []


def test_missing_database_is_safe_noop(tmp_path: Path) -> None:
    base = [
        base_result(
            "pages/students/2016/jan_holp/README.md",
            "Ján Holp",
        )
    ]

    result = expand_results_with_supplemental_sources(
        tmp_path / "missing.sqlite",
        "Ján Holp",
        base,
        published_only=True,
        limit=5,
    )

    assert result == base