from __future__ import annotations

import json
import sqlite3
from pathlib import Path
from typing import Any

import pytest

import scripts.rag_utils as rag_utils
from scripts.rag_document_expansion import (
    detect_section_intent,
    expand_results_with_exact_document_section,
    find_exact_student_document,
    load_best_document_section_chunk,
)


MAROS_PATH = (
    "pages/students/2016/"
    "maros_harahus/README.md"
)

TOMAS_PATH = (
    "pages/students/2016/"
    "tomas_kucharik/README.md"
)

DAVID_PATH = (
    "pages/students/2019/"
    "david_stromp/README.md"
)


def create_test_db(
    tmp_path: Path,
) -> Path:
    db_path = (
        tmp_path
        / "rag_document_expansion.sqlite"
    )

    with sqlite3.connect(
        db_path
    ) as conn:
        conn.executescript(
            """
            CREATE TABLE documents (
                id INTEGER PRIMARY KEY,
                path TEXT UNIQUE NOT NULL,
                title TEXT,
                author TEXT,
                published INTEGER,
                content_length INTEGER
                    NOT NULL DEFAULT 0,
                metadata_json TEXT
                    NOT NULL DEFAULT '{}'
            );

            CREATE TABLE chunks (
                id INTEGER PRIMARY KEY,
                chunk_id TEXT UNIQUE NOT NULL,
                document_path TEXT NOT NULL,
                title TEXT,
                author TEXT,
                published INTEGER,
                chunk_index INTEGER NOT NULL,
                heading_paths_json TEXT
                    NOT NULL DEFAULT '[]',
                text TEXT NOT NULL,
                text_length INTEGER
                    NOT NULL DEFAULT 0
            );
            """
        )

    return db_path


def insert_document(
    db_path: Path,
    *,
    path: str,
    title: str,
    author: str = "Daniel Hladek",
    published: bool = True,
) -> None:
    with sqlite3.connect(
        db_path
    ) as conn:
        conn.execute(
            """
            INSERT INTO documents (
                path,
                title,
                author,
                published
            )
            VALUES (?, ?, ?, ?)
            """,
            (
                path,
                title,
                author,
                (
                    1
                    if published
                    else 0
                ),
            ),
        )


def insert_chunk(
    db_path: Path,
    *,
    chunk_id: str,
    document_path: str,
    title: str,
    author: str,
    chunk_index: int,
    heading_paths: list[Any],
    text: str,
    published: bool = True,
) -> None:
    with sqlite3.connect(
        db_path
    ) as conn:
        conn.execute(
            """
            INSERT INTO chunks (
                chunk_id,
                document_path,
                title,
                author,
                published,
                chunk_index,
                heading_paths_json,
                text,
                text_length
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                chunk_id,
                document_path,
                title,
                author,
                (
                    1
                    if published
                    else 0
                ),
                chunk_index,
                json.dumps(
                    heading_paths,
                    ensure_ascii=False,
                ),
                text,
                len(
                    text
                ),
            ),
        )


def add_maros_document(
    db_path: Path,
) -> None:
    insert_document(
        db_path,
        path=MAROS_PATH,
        title="Maroš Harahus",
    )

    insert_chunk(
        db_path,
        chunk_id=(
            f"{MAROS_PATH}::chunk-2"
        ),
        document_path=(
            MAROS_PATH
        ),
        title="Maroš Harahus",
        author="Daniel Hladek",
        chunk_index=2,
        heading_paths=[
            [
                "Maroš Harahus",
                (
                    "Druhý rok "
                    "doktorandského štúdia"
                ),
            ],
        ],
        text=(
            "Ciele doktorandského štúdia. "
            "Obhájiť minimovku a "
            "publikovať článok."
        ),
    )

    insert_chunk(
        db_path,
        chunk_id=(
            f"{MAROS_PATH}::chunk-20"
        ),
        document_path=(
            MAROS_PATH
        ),
        title="Maroš Harahus",
        author="Daniel Hladek",
        chunk_index=20,
        heading_paths=[
            [
                "Maroš Harahus",
                "Diplomová práca 2021",
            ],
        ],
        text=(
            "Dokument: Maroš Harahus\n"
            "Sekcia: Diplomová práca 2021\n\n"
            "Názov diplomovej práce: "
            "Neurónová morfologická anotácia "
            "slovenského jazyka\n\n"
            "Zadanie diplomovej práce."
        ),
    )


def add_david_document(
    db_path: Path,
) -> None:
    insert_document(
        db_path,
        path=DAVID_PATH,
        title="Dávid Stromp",
        author="Maroš Harahus",
    )

    insert_chunk(
        db_path,
        chunk_id=(
            f"{DAVID_PATH}::chunk-0"
        ),
        document_path=(
            DAVID_PATH
        ),
        title="Dávid Stromp",
        author="Maroš Harahus",
        chunk_index=0,
        heading_paths=[
            [
                "Bakalárska práca 2022"
            ],
        ],
        text=(
            "Dokument: Dávid Stromp\n"
            "Autor dokumentu: Maroš Harahus\n"
            "Sekcia: Bakalárska práca 2022."
        ),
    )


def add_tomas_document(
    db_path: Path,
) -> None:
    insert_document(
        db_path,
        path=TOMAS_PATH,
        title="Tomáš Kuchárik",
    )

    insert_chunk(
        db_path,
        chunk_id=(
            f"{TOMAS_PATH}::chunk-0"
        ),
        document_path=(
            TOMAS_PATH
        ),
        title="Tomáš Kuchárik",
        author="Daniel Hladek",
        chunk_index=0,
        heading_paths=[
            [
                "Tomáš Kuchárik"
            ],
        ],
        text=(
            "Dokument: Tomáš Kuchárik\n"
            "(študent KPI)\n"
            "Súvisiace práce."
        ),
    )

    insert_chunk(
        db_path,
        chunk_id=(
            f"{TOMAS_PATH}::chunk-1"
        ),
        document_path=(
            TOMAS_PATH
        ),
        title="Tomáš Kuchárik",
        author="Daniel Hladek",
        chunk_index=1,
        heading_paths=[
            [
                "Tomáš Kuchárik",
                "Diplomová práca 2022",
            ],
        ],
        text=(
            "Dokument: Tomáš Kuchárik\n"
            "Sekcia: Diplomová práca 2022\n\n"
            "Názov: Tvorba korpusu otázok "
            "a odpovedí v slovenskom jazyku "
            "pomocou strojového prekladu\n\n"
            "Zadanie:"
        ),
    )

    insert_chunk(
        db_path,
        chunk_id=(
            f"{TOMAS_PATH}::chunk-5"
        ),
        document_path=(
            TOMAS_PATH
        ),
        title="Tomáš Kuchárik",
        author="Daniel Hladek",
        chunk_index=5,
        heading_paths=[
            [
                "Tomáš Kuchárik",
                "Diplomová práca 2021",
            ],
        ],
        text=(
            "Dokument: Tomáš Kuchárik\n"
            "Sekcia: Diplomová práca 2021\n\n"
            "Názov: Tvorba korpusu otázok "
            "a odpovedí v slovenskom jazyku "
            "pomocou crowdsourcingu."
        ),
    )


def retrieval_result(
    *,
    document_path: str,
    chunk_id: str,
    chunk_index: int,
    title: str,
    author: str,
    heading_paths: list[Any],
    text: str,
) -> dict[str, Any]:
    return {
        "chunk_id": chunk_id,
        "chunk_index": (
            chunk_index
        ),
        "document_path": (
            document_path
        ),
        "title": title,
        "author": author,
        "published": True,
        "heading_paths": (
            heading_paths
        ),
        "text": text,
        "source_url": (
            "https://zp.kemt.fei.tuke.sk/"
            + document_path.replace(
                "pages/",
                "",
            ).replace(
                "/README.md",
                "",
            )
        ),
        "match_strategy": (
            "all_terms"
        ),
        "fts_rank": 1,
        "vector_rank": 1,
        "vector_score": 0.90,
        "hybrid_score": 0.05,
        "context_expansion": {
            "strategy": (
                "section_lead"
            ),
            "applied": False,
            "primary_chunk_id": (
                chunk_id
            ),
            "primary_chunk_index": (
                chunk_index
            ),
            "lead_chunk_id": None,
            "lead_chunk_index": None,
        },
    }


@pytest.mark.parametrize(
    (
        "query",
        "expected",
    ),
    [
        (
            (
                "Maroš Harahus "
                "diplomová práca téma"
            ),
            "diploma_thesis",
        ),
        (
            (
                "Tomáš Kuchárik "
                "bakalárska práca"
            ),
            "bachelor_thesis",
        ),
        (
            (
                "Ján Test "
                "dizertačná práca"
            ),
            "doctoral_thesis",
        ),
        (
            (
                "Ján Test "
                "diplomový projekt"
            ),
            "diploma_project",
        ),
        (
            "Maroš Harahus",
            None,
        ),
    ],
)
def test_detect_section_intent(
    query: str,
    expected: str | None,
) -> None:
    assert (
        detect_section_intent(
            query
        )
        == expected
    )


def test_find_exact_student_document_uses_title_not_author(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    add_david_document(
        db_path
    )

    with sqlite3.connect(
        db_path
    ) as conn:
        conn.row_factory = (
            sqlite3.Row
        )

        document = (
            find_exact_student_document(
                conn,
                (
                    "Maroš Harahus "
                    "diplomová práca rok"
                ),
                published_only=True,
            )
        )

    assert (
        document
        is not None
    )

    assert (
        document[
            "document_path"
        ]
        == MAROS_PATH
    )

    assert (
        document[
            "title"
        ]
        == "Maroš Harahus"
    )


def test_load_best_chunk_q0020_selects_maros_diploma_section(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    with sqlite3.connect(
        db_path
    ) as conn:
        conn.row_factory = (
            sqlite3.Row
        )

        document = (
            find_exact_student_document(
                conn,
                (
                    "Maroš Harahus "
                    "diplomová práca téma"
                ),
                published_only=True,
            )
        )

        assert (
            document
            is not None
        )

        chunk = (
            load_best_document_section_chunk(
                conn,
                query=(
                    "Maroš Harahus "
                    "diplomová práca téma"
                ),
                document=document,
                intent="diploma_thesis",
                published_only=True,
            )
        )

    assert chunk is not None

    assert (
        chunk[
            "chunk_id"
        ]
        == (
            f"{MAROS_PATH}::chunk-20"
        )
    )

    assert (
        "Neurónová morfologická "
        "anotácia slovenského jazyka"
        in chunk[
            "text"
        ]
    )


def test_q0020_expands_wrong_chunk_in_correct_document(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    primary = retrieval_result(
        document_path=(
            MAROS_PATH
        ),
        chunk_id=(
            f"{MAROS_PATH}::chunk-2"
        ),
        chunk_index=2,
        title="Maroš Harahus",
        author="Daniel Hladek",
        heading_paths=[
            [
                "Maroš Harahus",
                (
                    "Druhý rok "
                    "doktorandského štúdia"
                ),
            ],
        ],
        text=(
            "Ciele doktorandského štúdia."
        ),
    )

    expanded = (
        expand_results_with_exact_document_section(
            db_path,
            (
                "Maroš Harahus "
                "diplomová práca téma"
            ),
            [
                primary
            ],
            published_only=True,
            limit=5,
        )
    )

    assert (
        len(
            expanded
        )
        == 1
    )

    item = expanded[
        0
    ]

    assert (
        item[
            "chunk_id"
        ]
        == primary[
            "chunk_id"
        ]
    )

    assert (
        "Neurónová morfologická "
        "anotácia slovenského jazyka"
        in item[
            "exact_document_text"
        ]
    )

    assert (
        item[
            "document_expansion"
        ][
            "applied"
        ]
        is True
    )

    assert (
        item[
            "document_expansion"
        ][
            "added_source"
        ]
        is False
    )


def test_q0021_adds_exact_maros_document_before_author_document(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    add_david_document(
        db_path
    )

    david_result = (
        retrieval_result(
            document_path=(
                DAVID_PATH
            ),
            chunk_id=(
                f"{DAVID_PATH}::chunk-0"
            ),
            chunk_index=0,
            title="Dávid Stromp",
            author="Maroš Harahus",
            heading_paths=[
                [
                    "Bakalárska práca 2022"
                ],
            ],
            text=(
                "Autor dokumentu: "
                "Maroš Harahus."
            ),
        )
    )

    expanded = (
        expand_results_with_exact_document_section(
            db_path,
            (
                "Maroš Harahus "
                "diplomová práca rok"
            ),
            [
                david_result
            ],
            published_only=True,
            limit=5,
        )
    )

    assert (
        len(
            expanded
        )
        == 2
    )

    assert (
        expanded[
            0
        ][
            "document_path"
        ]
        == MAROS_PATH
    )

    assert (
        expanded[
            0
        ][
            "title"
        ]
        == "Maroš Harahus"
    )

    assert (
        "Diplomová práca 2021"
        in expanded[
            0
        ][
            "text"
        ]
    )

    assert (
        expanded[
            0
        ][
            "document_expansion"
        ][
            "added_source"
        ]
        is True
    )

    assert (
        expanded[
            1
        ][
            "document_path"
        ]
        == DAVID_PATH
    )


def test_q0028_selects_answer_chunk_in_tomas_document(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_tomas_document(
        db_path
    )

    primary = retrieval_result(
        document_path=(
            TOMAS_PATH
        ),
        chunk_id=(
            f"{TOMAS_PATH}::chunk-0"
        ),
        chunk_index=0,
        title="Tomáš Kuchárik",
        author="Daniel Hladek",
        heading_paths=[
            [
                "Tomáš Kuchárik"
            ],
        ],
        text=(
            "Dokument: Tomáš Kuchárik\n"
            "Súvisiace práce."
        ),
    )

    expanded = (
        expand_results_with_exact_document_section(
            db_path,
            (
                "Tomáš Kuchárik "
                "diplomová práca"
            ),
            [
                primary
            ],
            published_only=True,
            limit=5,
        )
    )

    assert (
        len(
            expanded
        )
        == 1
    )

    assert (
        "Tvorba korpusu otázok "
        "a odpovedí v slovenskom jazyku "
        "pomocou strojového prekladu"
        in expanded[
            0
        ][
            "exact_document_text"
        ]
    )


def test_latest_work_section_is_preferred_without_year(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_tomas_document(
        db_path
    )

    with sqlite3.connect(
        db_path
    ) as conn:
        conn.row_factory = (
            sqlite3.Row
        )

        document = (
            find_exact_student_document(
                conn,
                (
                    "Tomáš Kuchárik "
                    "diplomová práca"
                ),
                published_only=True,
            )
        )

        assert document is not None

        chunk = (
            load_best_document_section_chunk(
                conn,
                query=(
                    "Tomáš Kuchárik "
                    "diplomová práca"
                ),
                document=document,
                intent="diploma_thesis",
                published_only=True,
            )
        )

    assert chunk is not None

    assert (
        chunk[
            "chunk_index"
        ]
        == 1
    )

    assert (
        "strojového prekladu"
        in chunk[
            "text"
        ]
    )


def test_explicit_year_selects_requested_work_section(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_tomas_document(
        db_path
    )

    with sqlite3.connect(
        db_path
    ) as conn:
        conn.row_factory = (
            sqlite3.Row
        )

        document = (
            find_exact_student_document(
                conn,
                (
                    "Tomáš Kuchárik "
                    "diplomová práca 2021"
                ),
                published_only=True,
            )
        )

        assert document is not None

        chunk = (
            load_best_document_section_chunk(
                conn,
                query=(
                    "Tomáš Kuchárik "
                    "diplomová práca 2021"
                ),
                document=document,
                intent="diploma_thesis",
                published_only=True,
            )
        )

    assert chunk is not None

    assert (
        chunk[
            "chunk_index"
        ]
        == 5
    )

    assert (
        "crowdsourcingu"
        in chunk[
            "text"
        ]
    )


def test_no_exact_title_leaves_results_unchanged(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    original = retrieval_result(
        document_path=(
            DAVID_PATH
        ),
        chunk_id=(
            f"{DAVID_PATH}::chunk-0"
        ),
        chunk_index=0,
        title="Dávid Stromp",
        author="Maroš Harahus",
        heading_paths=[
            [
                "Bakalárska práca 2022"
            ],
        ],
        text="Test",
    )

    expanded = (
        expand_results_with_exact_document_section(
            db_path,
            (
                "Neznámy Študent "
                "diplomová práca"
            ),
            [
                original
            ],
            published_only=True,
            limit=5,
        )
    )

    assert (
        expanded
        == [
            original
        ]
    )


def test_query_without_section_intent_leaves_results_unchanged(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    original = retrieval_result(
        document_path=(
            MAROS_PATH
        ),
        chunk_id=(
            f"{MAROS_PATH}::chunk-2"
        ),
        chunk_index=2,
        title="Maroš Harahus",
        author="Daniel Hladek",
        heading_paths=[
            [
                "Maroš Harahus"
            ],
        ],
        text="Test",
    )

    expanded = (
        expand_results_with_exact_document_section(
            db_path,
            "Maroš Harahus",
            [
                original
            ],
            published_only=True,
            limit=5,
        )
    )

    assert (
        expanded
        == [
            original
        ]
    )


def test_published_only_does_not_use_unpublished_document(
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    insert_document(
        db_path,
        path=MAROS_PATH,
        title="Maroš Harahus",
        published=False,
    )

    with sqlite3.connect(
        db_path
    ) as conn:
        conn.row_factory = (
            sqlite3.Row
        )

        document = (
            find_exact_student_document(
                conn,
                (
                    "Maroš Harahus "
                    "diplomová práca"
                ),
                published_only=True,
            )
        )

    assert (
        document
        is None
    )


def test_build_source_text_includes_exact_document_section() -> None:
    result = retrieval_result(
        document_path=(
            TOMAS_PATH
        ),
        chunk_id=(
            f"{TOMAS_PATH}::chunk-0"
        ),
        chunk_index=0,
        title="Tomáš Kuchárik",
        author="Daniel Hladek",
        heading_paths=[
            [
                "Tomáš Kuchárik"
            ],
        ],
        text=(
            "Primárny retrieval úsek."
        ),
    )

    result[
        "exact_document_text"
    ] = (
        "Diplomová práca 2022\n"
        "Názov: Tvorba korpusu otázok "
        "a odpovedí v slovenskom jazyku "
        "pomocou strojového prekladu"
    )

    text = (
        rag_utils.build_source_text(
            result
        )
    )

    assert (
        "RELEVANTNÁ SEKCIA PRESNE "
        "ZHODNÉHO DOKUMENTU"
        in text
    )

    assert (
        "strojového prekladu"
        in text
    )

    assert (
        "Primárny retrieval úsek."
        in text
    )


def test_build_rag_context_q0028_regression(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_tomas_document(
        db_path
    )

    primary = retrieval_result(
        document_path=(
            TOMAS_PATH
        ),
        chunk_id=(
            f"{TOMAS_PATH}::chunk-0"
        ),
        chunk_index=0,
        title="Tomáš Kuchárik",
        author="Daniel Hladek",
        heading_paths=[
            [
                "Tomáš Kuchárik"
            ],
        ],
        text=(
            "Dokument: Tomáš Kuchárik\n"
            "Súvisiace práce."
        ),
    )

    def fake_search_database(
        db_path_arg: Path,
        query: str,
        limit: int,
        *,
        published_only: bool,
        max_per_document: int,
    ) -> dict[str, Any]:
        assert (
            db_path_arg
            == db_path
        )

        return {
            "engine": (
                "hybrid_fts5_embeddings"
            ),
            "strategies": [
                "all_terms"
            ],
            "results": [
                primary
            ],
        }

    monkeypatch.setattr(
        rag_utils,
        "search_database",
        fake_search_database,
    )

    response = (
        rag_utils.build_rag_context(
            db_path,
            (
                "Tomáš Kuchárik "
                "diplomová práca"
            ),
            limit=5,
            published_only=True,
            max_per_document=1,
        )
    )

    assert (
        response[
            "source_count"
        ]
        == 1
    )

    assert (
        "Tvorba korpusu otázok "
        "a odpovedí v slovenskom jazyku "
        "pomocou strojového prekladu"
        in response[
            "context"
        ]
    )

    assert (
        response[
            "sources"
        ][0][
            "source_url"
        ]
        == (
            "https://zp.kemt.fei.tuke.sk/"
            "students/2016/tomas_kucharik"
        )
    )


def test_build_rag_context_q0021_regression(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    db_path = (
        create_test_db(
            tmp_path
        )
    )

    add_maros_document(
        db_path
    )

    add_david_document(
        db_path
    )

    david_result = (
        retrieval_result(
            document_path=(
                DAVID_PATH
            ),
            chunk_id=(
                f"{DAVID_PATH}::chunk-0"
            ),
            chunk_index=0,
            title="Dávid Stromp",
            author="Maroš Harahus",
            heading_paths=[
                [
                    "Bakalárska práca 2022"
                ],
            ],
            text=(
                "Dokument: Dávid Stromp\n"
                "Autor dokumentu: "
                "Maroš Harahus."
            ),
        )
    )

    def fake_search_database(
        db_path_arg: Path,
        query: str,
        limit: int,
        *,
        published_only: bool,
        max_per_document: int,
    ) -> dict[str, Any]:
        return {
            "engine": (
                "hybrid_fts5_embeddings"
            ),
            "strategies": [
                "prefix_terms"
            ],
            "results": [
                david_result
            ],
        }

    monkeypatch.setattr(
        rag_utils,
        "search_database",
        fake_search_database,
    )

    response = (
        rag_utils.build_rag_context(
            db_path,
            (
                "Maroš Harahus "
                "diplomová práca rok"
            ),
            limit=5,
            published_only=True,
            max_per_document=1,
        )
    )

    assert (
        response[
            "sources"
        ][0][
            "title"
        ]
        == "Maroš Harahus"
    )

    assert (
        "Diplomová práca 2021"
        in response[
            "context"
        ]
    )

    assert (
        "Neurónová morfologická "
        "anotácia slovenského jazyka"
        in response[
            "context"
        ]
    )