from __future__ import annotations

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

from scripts.search_utils import (
    make_source_url,
    normalize_for_compare,
)


YEAR_RE = re.compile(
    r"\b(?:19|20)\d{2}\b"
)


SECTION_INTENTS: dict[
    str,
    tuple[str, ...],
] = {
    "diploma_thesis": (
        "diplomov",
        "prac",
    ),
    "bachelor_thesis": (
        "bakalarsk",
        "prac",
    ),
    "doctoral_thesis": (
        "dizertacn",
        "prac",
    ),
    "diploma_project": (
        "diplomov",
        "projekt",
    ),
    "team_project": (
        "timov",
        "projekt",
    ),
}


WORK_INTENTS = {
    "diploma_thesis",
    "bachelor_thesis",
    "doctoral_thesis",
}


QUERY_STOPWORDS = {
    "a",
    "aj",
    "aka",
    "ake",
    "aky",
    "je",
    "k",
    "kde",
    "ktora",
    "ktore",
    "ktory",
    "ma",
    "mal",
    "mala",
    "malo",
    "na",
    "o",
    "osoba",
    "osobe",
    "osoby",
    "pri",
    "pre",
    "sa",
    "s",
    "so",
    "su",
    "ta",
    "ten",
    "to",
    "u",
    "v",
    "vo",
    "z",
    "za",
}


def normalized_tokens(
    value: str,
) -> list[str]:
    normalized = (
        normalize_for_compare(
            value
        )
    )

    if not normalized:
        return []

    return normalized.split()


def has_prefix(
    tokens: list[str],
    prefix: str,
) -> bool:
    return any(
        token.startswith(
            prefix
        )
        for token in tokens
    )


def detect_section_intent(
    query: str,
) -> str | None:
    tokens = normalized_tokens(
        query
    )

    if not tokens:
        return None

    # Diplomový projekt musí byť pred
    # diplomovou prácou, pretože oba
    # začínajú tokenom "diplomov".
    ordered_intents = (
        "diploma_project",
        "team_project",
        "diploma_thesis",
        "bachelor_thesis",
        "doctoral_thesis",
    )

    for intent in ordered_intents:
        prefixes = (
            SECTION_INTENTS[
                intent
            ]
        )

        if all(
            has_prefix(
                tokens,
                prefix,
            )
            for prefix in prefixes
        ):
            return intent

    return None


def extract_years(
    value: str,
) -> list[int]:
    return [
        int(
            match
        )
        for match in YEAR_RE.findall(
            value
        )
    ]


def parse_heading_paths(
    value: Any,
) -> list[Any]:
    if isinstance(
        value,
        list,
    ):
        return value

    if not value:
        return []

    try:
        parsed = json.loads(
            str(
                value
            )
        )

    except json.JSONDecodeError:
        return []

    if not isinstance(
        parsed,
        list,
    ):
        return []

    return parsed


def flatten_heading_paths(
    heading_paths: list[Any],
) -> str:
    parts: list[str] = []

    for item in heading_paths:
        if isinstance(
            item,
            str,
        ):
            value = (
                item.strip()
            )

            if value:
                parts.append(
                    value
                )

            continue

        if isinstance(
            item,
            (
                list,
                tuple,
            ),
        ):
            for part in item:
                value = str(
                    part
                ).strip()

                if value:
                    parts.append(
                        value
                    )

    return " ".join(
        parts
    )


def section_matches_intent(
    value: str,
    intent: str,
) -> bool:
    prefixes = (
        SECTION_INTENTS.get(
            intent
        )
    )

    if not prefixes:
        return False

    tokens = normalized_tokens(
        value
    )

    return all(
        has_prefix(
            tokens,
            prefix,
        )
        for prefix in prefixes
    )


def is_primary_student_document(
    document_path: str,
) -> bool:
    parts = Path(
        document_path
    ).parts

    return (
        len(
            parts
        )
        == 5
        and parts[0]
        == "pages"
        and parts[1]
        == "students"
        and parts[-1]
        == "README.md"
    )


def title_occurs_in_query(
    title: str,
    query: str,
) -> bool:
    normalized_title = (
        normalize_for_compare(
            title
        )
    )

    normalized_query = (
        normalize_for_compare(
            query
        )
    )

    if (
        not normalized_title
        or not normalized_query
    ):
        return False

    # Jednoslovné názvy sú príliš široké
    # na bezpečné exact-document rozšírenie.
    if len(
        normalized_title.split()
    ) < 2:
        return False

    return (
        f" {normalized_title} "
        in f" {normalized_query} "
    )


def sqlite_table_exists(
    conn: sqlite3.Connection,
    table_name: str,
) -> bool:
    row = conn.execute(
        """
        SELECT 1
        FROM sqlite_master
        WHERE type = 'table'
          AND name = ?
        LIMIT 1
        """,
        (
            table_name,
        ),
    ).fetchone()

    return row is not None


def find_exact_student_document(
    conn: sqlite3.Connection,
    query: str,
    *,
    published_only: bool,
) -> dict[str, Any] | None:
    rows = conn.execute(
        """
        SELECT
            path,
            title,
            author,
            published
        FROM documents
        WHERE (
            ? = 0
            OR published = 1
        )
        ORDER BY
            LENGTH(title) DESC,
            path ASC
        """,
        (
            1
            if published_only
            else 0,
        ),
    ).fetchall()

    candidates: list[
        dict[str, Any]
    ] = []

    for row in rows:
        document_path = str(
            row[
                "path"
            ]
            or ""
        ).strip()

        title = str(
            row[
                "title"
            ]
            or ""
        ).strip()

        if (
            not document_path
            or not title
        ):
            continue

        if not is_primary_student_document(
            document_path
        ):
            continue

        if not title_occurs_in_query(
            title,
            query,
        ):
            continue

        candidates.append(
            {
                "document_path": (
                    document_path
                ),
                "title": title,
                "author": row[
                    "author"
                ],
                "published": (
                    row[
                        "published"
                    ]
                ),
            }
        )

    if not candidates:
        return None

    candidates.sort(
        key=lambda item: (
            -len(
                normalized_tokens(
                    str(
                        item[
                            "title"
                        ]
                    )
                )
            ),
            -len(
                str(
                    item[
                        "title"
                    ]
                )
            ),
            str(
                item[
                    "document_path"
                ]
            ),
        )
    )

    return candidates[
        0
    ]


def query_content_tokens(
    query: str,
    title: str,
) -> list[str]:
    query_tokens = (
        normalized_tokens(
            query
        )
    )

    title_tokens = set(
        normalized_tokens(
            title
        )
    )

    result: list[str] = []

    for token in query_tokens:
        if token in title_tokens:
            continue

        if token in QUERY_STOPWORDS:
            continue

        if len(token) < 3:
            continue

        result.append(
            token
        )

    return result


def local_chunk_score(
    *,
    query: str,
    title: str,
    intent: str,
    heading_paths: list[Any],
    text: str,
    chunk_index: int,
) -> tuple[
    int,
    int,
    int,
]:
    heading_text = (
        flatten_heading_paths(
            heading_paths
        )
    )

    heading_normalized = (
        normalize_for_compare(
            heading_text
        )
    )

    text_normalized = (
        normalize_for_compare(
            text
        )
    )

    score = 0

    heading_matches = (
        section_matches_intent(
            heading_text,
            intent,
        )
    )

    text_matches = (
        section_matches_intent(
            text,
            intent,
        )
    )

    if heading_matches:
        score += 1000

    if text_matches:
        score += 500

    # Ak chunk vôbec nepatrí k požadovanému
    # typu sekcie, nemá byť lokálnym víťazom.
    if (
        not heading_matches
        and not text_matches
    ):
        return (
            -1,
            -1,
            -chunk_index,
        )

    query_years = extract_years(
        query
    )

    heading_years = extract_years(
        heading_text
    )

    text_years = extract_years(
        text
    )

    if query_years:
        if any(
            year in heading_years
            for year in query_years
        ):
            score += 1500

        elif any(
            year in text_years
            for year in query_years
        ):
            score += 750

        else:
            score -= 500

    content_tokens = (
        query_content_tokens(
            query,
            title,
        )
    )

    for token in content_tokens:
        if token in heading_normalized:
            score += 30

        if token in text_normalized:
            score += 8

    if intent in WORK_INTENTS:
        if (
            "nazov diplomovej prace"
            in text_normalized
            or "nazov bakalarskej prace"
            in text_normalized
            or "nazov dizertacnej prace"
            in text_normalized
        ):
            score += 120

        elif "nazov" in text_normalized:
            score += 80

        if "tema" in text_normalized:
            score += 50

    if has_prefix(
        normalized_tokens(
            query
        ),
        "rok",
    ):
        if heading_years:
            score += 100

    # Ak rok nie je v dopyte a dokument má
    # viac rovnakých typov prác, preferujeme
    # najnovšiu sekciu.
    latest_year = (
        max(
            heading_years
        )
        if heading_years
        else -1
    )

    return (
        score,
        latest_year,
        -chunk_index,
    )


def load_best_document_section_chunk(
    conn: sqlite3.Connection,
    *,
    query: str,
    document: dict[str, Any],
    intent: str,
    published_only: bool,
) -> dict[str, Any] | None:
    document_path = str(
        document[
            "document_path"
        ]
    )

    rows = conn.execute(
        """
        SELECT
            chunk_id,
            document_path,
            title,
            author,
            published,
            chunk_index,
            heading_paths_json,
            text
        FROM chunks
        WHERE document_path = ?
          AND (
              ? = 0
              OR published = 1
          )
        ORDER BY
            chunk_index ASC,
            id ASC
        """,
        (
            document_path,
            (
                1
                if published_only
                else 0
            ),
        ),
    ).fetchall()

    best_item: dict[
        str,
        Any,
    ] | None = None

    best_score: tuple[
        int,
        int,
        int,
    ] | None = None

    for row in rows:
        heading_paths = (
            parse_heading_paths(
                row[
                    "heading_paths_json"
                ]
            )
        )

        text = str(
            row[
                "text"
            ]
            or ""
        ).strip()

        if not text:
            continue

        chunk_index = int(
            row[
                "chunk_index"
            ]
        )

        score = local_chunk_score(
            query=query,
            title=str(
                document[
                    "title"
                ]
            ),
            intent=intent,
            heading_paths=(
                heading_paths
            ),
            text=text,
            chunk_index=(
                chunk_index
            ),
        )

        if score[
            0
        ] < 0:
            continue

        if (
            best_score is not None
            and score <= best_score
        ):
            continue

        best_score = score

        best_item = {
            "chunk_id": str(
                row[
                    "chunk_id"
                ]
            ),
            "document_path": (
                document_path
            ),
            "title": (
                row[
                    "title"
                ]
                or document[
                    "title"
                ]
            ),
            "author": (
                row[
                    "author"
                ]
                or document.get(
                    "author"
                )
            ),
            "published": (
                bool(
                    row[
                        "published"
                    ]
                )
                if row[
                    "published"
                ]
                is not None
                else None
            ),
            "chunk_index": (
                chunk_index
            ),
            "heading_paths": (
                heading_paths
            ),
            "text": text,
            "source_url": (
                make_source_url(
                    document_path
                )
            ),
        }

    return best_item


def exact_document_expansion_metadata(
    *,
    applied: bool,
    document_path: str | None,
    chunk_id: str | None,
    chunk_index: int | None,
    added_source: bool,
) -> dict[str, Any]:
    return {
        "strategy": (
            "exact_document_section"
        ),
        "applied": applied,
        "document_path": (
            document_path
        ),
        "chunk_id": (
            chunk_id
        ),
        "chunk_index": (
            chunk_index
        ),
        "added_source": (
            added_source
        ),
    }


def expand_results_with_exact_document_section(
    db_path: Path,
    query: str,
    results: list[
        dict[str, Any]
    ],
    *,
    published_only: bool,
    limit: int,
) -> list[dict[str, Any]]:
    if not results:
        base_results: list[
            dict[str, Any]
        ] = []

    else:
        base_results = [
            dict(
                result
            )
            for result in results
        ]

    intent = detect_section_intent(
        query
    )

    if intent is None:
        return base_results

    if not db_path.exists():
        return base_results

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

        conn.execute(
            "PRAGMA query_only = ON"
        )

        # Niektoré unit testy používajú zámerne
        # minimalistickú SQLite databázu iba
        # s tabuľkou chunks.
        #
        # Exact-document expansion je voliteľná
        # RAG-only nadstavba. Ak potrebná schéma
        # nie je dostupná, musí zostať bezpečný
        # pôvodný výsledok retrievalu.
        if (
            not sqlite_table_exists(
                conn,
                "documents",
            )
            or not sqlite_table_exists(
                conn,
                "chunks",
            )
        ):
            return base_results

        document = (
            find_exact_student_document(
                conn,
                query,
                published_only=(
                    published_only
                ),
            )
        )

        if document is None:
            return base_results

        local_chunk = (
            load_best_document_section_chunk(
                conn,
                query=query,
                document=document,
                intent=intent,
                published_only=(
                    published_only
                ),
            )
        )

    if local_chunk is None:
        return base_results

    document_path = str(
        document[
            "document_path"
        ]
    )

    existing_index: int | None = None

    for index, result in enumerate(
        base_results
    ):
        if str(
            result.get(
                "document_path"
            )
            or ""
        ) == document_path:
            existing_index = (
                index
            )

            break

    if existing_index is not None:
        item = dict(
            base_results[
                existing_index
            ]
        )

        local_chunk_id = str(
            local_chunk[
                "chunk_id"
            ]
        )

        primary_chunk_id = str(
            item.get(
                "chunk_id"
            )
            or ""
        )

        if (
            local_chunk_id
            != primary_chunk_id
        ):
            item[
                "exact_document_text"
            ] = local_chunk[
                "text"
            ]

            item[
                "exact_document_heading_paths"
            ] = local_chunk[
                "heading_paths"
            ]

        item[
            "document_expansion"
        ] = (
            exact_document_expansion_metadata(
                applied=(
                    local_chunk_id
                    != primary_chunk_id
                ),
                document_path=(
                    document_path
                ),
                chunk_id=(
                    local_chunk_id
                ),
                chunk_index=int(
                    local_chunk[
                        "chunk_index"
                    ]
                ),
                added_source=False,
            )
        )

        del base_results[
            existing_index
        ]

        # Presne zhodný dokument má byť
        # v RAG kontexte pred dokumentmi,
        # kde je meno iba v author poli.
        base_results.insert(
            0,
            item,
        )

        return (
            base_results[
                :limit
            ]
            if limit > 0
            else base_results
        )

    synthetic = {
        **local_chunk,
        "match_strategy": (
            "exact_document_section"
        ),
        "fts_rank": None,
        "vector_rank": None,
        "vector_score": None,
        "hybrid_score": None,
        "document_expansion": (
            exact_document_expansion_metadata(
                applied=True,
                document_path=(
                    document_path
                ),
                chunk_id=str(
                    local_chunk[
                        "chunk_id"
                    ]
                ),
                chunk_index=int(
                    local_chunk[
                        "chunk_index"
                    ]
                ),
                added_source=True,
            )
        ),
    }

    expanded = [
        synthetic,
        *base_results,
    ]

    if limit > 0:
        return expanded[
            :limit
        ]

    return expanded