from __future__ import annotations

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

import pytest

import tools.rag.rag_utils as rag_utils
from tools.rag.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"]