from __future__ import annotations

import sqlite3
from pathlib import Path

from scripts.rag_query_evidence import (
    build_focus_excerpt,
    expand_results_with_query_evidence,
    find_best_query_evidence_chunk,
)


def make_db(
    tmp_path: Path,
) -> Path:
    db = (
        tmp_path
        / "evidence.sqlite"
    )

    with sqlite3.connect(
        db
    ) as conn:
        conn.execute(
            """
            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
            )
            """
        )

    return db


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


def result(
    path: str,
    index: int,
    title: str,
    text: str,
) -> dict:
    return {
        "chunk_id": (
            f"{path}::chunk-{index}"
        ),
        "document_path": path,
        "title": title,
        "author": "Daniel Hladek",
        "published": True,
        "chunk_index": index,
        "heading_paths": [],
        "text": text,
        "source_url": (
            "https://example.test/student"
        ),
        "match_strategy": "hybrid",
        "fts_rank": 1,
        "vector_rank": 1,
        "vector_score": 0.9,
        "hybrid_score": 0.03,
    }


def test_reverse_title_selects_tomas_2022_chunk(
    tmp_path: Path,
) -> None:
    db = make_db(
        tmp_path
    )

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

    exact_title = (
        "Tvorba korpusu otázok a odpovedí "
        "v slovenskom jazyku pomocou "
        "strojového prekladu"
    )

    add_chunk(
        db,
        path=path,
        index=1,
        title="Tomáš Kuchárik",
        heading=(
            '[["Tomáš Kuchárik", '
            '"Diplomová práca 2022"]]'
        ),
        text=(
            "## Diplomová práca 2022\n\n"
            f"Názov: {exact_title}"
        ),
    )

    add_chunk(
        db,
        path=path,
        index=5,
        title="Tomáš Kuchárik",
        heading=(
            '[["Tomáš Kuchárik", '
            '"Diplomová práca 2021"]]'
        ),
        text=(
            "## Diplomová práca 2021\n\n"
            "Názov: Tvorba korpusu otázok "
            "a odpovedí v slovenskom jazyku "
            "pomocou crowdsourcingu"
        ),
    )

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

        best = (
            find_best_query_evidence_chunk(
                conn,
                path,
                f"{exact_title} autor",
            )
        )

    assert best is not None

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

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


def test_expansion_preserves_retrieval_chunk_and_adds_2022_evidence(
    tmp_path: Path,
) -> None:
    db = make_db(
        tmp_path
    )

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

    exact_title = (
        "Tvorba korpusu otázok a odpovedí "
        "v slovenskom jazyku pomocou "
        "strojového prekladu"
    )

    add_chunk(
        db,
        path=path,
        index=1,
        title="Tomáš Kuchárik",
        heading=(
            '[["Tomáš Kuchárik", '
            '"Diplomová práca 2022"]]'
        ),
        text=(
            f"Názov: {exact_title}"
        ),
    )

    add_chunk(
        db,
        path=path,
        index=5,
        title="Tomáš Kuchárik",
        heading=(
            '[["Tomáš Kuchárik", '
            '"Diplomová práca 2021"]]'
        ),
        text=(
            "Názov: Tvorba korpusu "
            "pomocou crowdsourcingu"
        ),
    )

    original = result(
        path,
        5,
        "Tomáš Kuchárik",
        "crowdsourcing",
    )

    expanded = (
        expand_results_with_query_evidence(
            db,
            exact_title,
            [
                original
            ],
        )
    )

    item = expanded[
        0
    ]

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

    assert (
        item[
            "hybrid_score"
        ]
        == original[
            "hybrid_score"
        ]
    )

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

    assert (
        item[
            "query_evidence"
        ][
            "evidence_chunk_index"
        ]
        == 1
    )

    assert (
        "strojového prekladu"
        in item[
            "query_focus_text"
        ]
    )


def test_us_steel_focus_contains_gnn_task() -> None:
    text = (
        "Stretnutie 1.10.\n"
        "Stav:\n"
        "- Štúdium základov neurónových sietí\n"
        "- Úvodné stretnutie s US Steel\n"
        "Úlohy:\n"
        "- Vypracovať prehľad aktuálnych "
        "metód grafových neurónových sietí\n"
        "- Nájsť a vyskúšať toolkit na GNN.\n"
        "- Naštudovať dáta z US Steel."
    )

    excerpt = (
        build_focus_excerpt(
            (
                "Maroš Harahus US Steel "
                "metódy študovať"
            ),
            text,
        )
    )

    assert (
        "US Steel"
        in excerpt
    )

    assert (
        "grafových neurónových sietí"
        in excerpt
    )

    assert (
        "toolkit na GNN"
        in excerpt
    )


def test_maros_own_document_gets_gnn_evidence(
    tmp_path: Path,
) -> None:
    db = make_db(
        tmp_path
    )

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

    text = (
        "Úlohy:\n"
        "- Vypracovať prehľad aktuálnych "
        "metód grafových neurónových sietí\n"
        "- Nájsť a vyskúšať toolkit na GNN."
    )

    add_chunk(
        db,
        path=path,
        index=19,
        title="Maroš Harahus",
        heading=(
            '[["Maroš Harahus", '
            '"Prvý ročník PhD štúdia"]]'
        ),
        text=text,
    )

    expanded = (
        expand_results_with_query_evidence(
            db,
            (
                "Maroš Harahus "
                "grafové neurónové siete"
            ),
            [
                result(
                    path,
                    19,
                    "Maroš Harahus",
                    text,
                )
            ],
        )
    )

    assert (
        expanded[
            0
        ][
            "query_evidence"
        ][
            "applied"
        ]
        is True
    )

    assert (
        "grafových neurónových sietí"
        in expanded[
            0
        ][
            "query_focus_text"
        ]
    )


def test_published_only_ignores_unpublished_better_chunk(
    tmp_path: Path,
) -> None:
    db = make_db(
        tmp_path
    )

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

    query = (
        "Tvorba korpusu otázok a odpovedí "
        "pomocou strojového prekladu"
    )

    add_chunk(
        db,
        path=path,
        index=1,
        title="Tomáš Kuchárik",
        heading=(
            '[["Diplomová práca 2022"]]'
        ),
        text=(
            "strojový preklad"
        ),
    )

    add_chunk(
        db,
        path=path,
        index=2,
        title="Tomáš Kuchárik",
        heading=(
            '[["Diplomová práca 2022"]]'
        ),
        text=query,
        published=False,
    )

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

        best = (
            find_best_query_evidence_chunk(
                conn,
                path,
                query,
                published_only=True,
            )
        )

    assert best is not None

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


def test_missing_chunks_table_is_safe_noop(
    tmp_path: Path,
) -> None:
    db = (
        tmp_path
        / "empty.sqlite"
    )

    with sqlite3.connect(
        db
    ):
        pass

    original = result(
        (
            "pages/students/2016/"
            "test/README.md"
        ),
        0,
        "Test",
        "text",
    )

    assert (
        expand_results_with_query_evidence(
            db,
            "nejaký dotaz",
            [
                original
            ],
        )
        == [
            original
        ]
    )