from __future__ import annotations

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

from tools.rag.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 COALESCE(published, 1) = 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 COALESCE(published, 1) = 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