from __future__ import annotations

import json
import re
import sqlite3
import unicodedata
from collections import defaultdict
from typing import Any

from tools.rag.embedding_utils import (
    blob_to_vector,
    cosine_similarity,
    embed_query,
    model_profile,
)

WORD_RE = re.compile(
    r"[^\W_]+",
    re.UNICODE,
)

QUOTED_PHRASE_RE = re.compile(
    r'[„“”"]([^„“”"]+)[„“”"]',
    re.UNICODE,
)


BM25_WEIGHTS = (
    0.0,
    10.0,
    7.0,
    5.0,
    8.0,
    4.0,
    1.0,
)

BM25_SQL = ", ".join(str(value) for value in BM25_WEIGHTS)


DEFAULT_CANDIDATE_MULTIPLIER = 8
MIN_CANDIDATES = 50
MIN_STEM_PREFIX_LENGTH = 5

RRF_K = 60
FTS_RRF_WEIGHT = 1.0
VECTOR_RRF_WEIGHT = 1.5
ANY_TERM_RRF_WEIGHT = 0.25

# Ak FTS kandidát obsahuje veľmi silný
# lexikálny signál, napríklad meno osoby
# alebo presný názov práce, nesmie byť
# potlačený iba embeddingovým výsledkom.
LEXICAL_ANCHOR_MIN_BONUS = 4.5
LEXICAL_ANCHOR_RRF_WEIGHT = 2.0


STRATEGY_PRIORITY = {
    "all_terms": 3,
    "prefix_terms": 2,
    "any_term": 1,
}


def normalize_for_compare(
    text: str,
) -> str:
    text = unicodedata.normalize(
        "NFKD",
        text.casefold(),
    )

    text = "".join(character for character in text if not unicodedata.combining(character))

    return " ".join(WORD_RE.findall(text))


def query_tokens(
    query: str,
) -> list[str]:
    tokens: list[str] = []
    seen: set[str] = set()

    for token in WORD_RE.findall(query):
        normalized = normalize_for_compare(token)

        if not normalized:
            continue

        if normalized in seen:
            continue

        tokens.append(token)

        seen.add(normalized)

    return tokens


def quoted_query_phrases(
    query: str,
) -> list[str]:
    phrases: list[str] = []
    seen: set[str] = set()

    for match in QUOTED_PHRASE_RE.finditer(query):
        phrase = normalize_for_compare(match.group(1))

        if not phrase:
            continue

        # Jednoslovné výrazy ako "RAG"
        # alebo "chatbot" sú príliš široké
        # na silný lexikálny anchor.
        if len(phrase.split()) < 2:
            continue

        if phrase in seen:
            continue

        phrases.append(phrase)

        seen.add(phrase)

    return phrases


def quote_fts_token(
    token: str,
    *,
    use_prefix: bool = True,
    shorten: bool = False,
) -> str:
    value = token

    if shorten and len(value) > MIN_STEM_PREFIX_LENGTH:
        value = value[:MIN_STEM_PREFIX_LENGTH]

    escaped = value.replace(
        '"',
        '""',
    )

    suffix = "*" if (use_prefix and len(value) >= 4) else ""

    return f'"{escaped}"{suffix}'


def build_match_queries(
    query: str,
) -> list[tuple[str, str]]:
    tokens = query_tokens(query)

    if not tokens:
        return []

    full_terms = [quote_fts_token(token) for token in tokens]

    all_terms_query = " AND ".join(full_terms)

    queries = [
        (
            "all_terms",
            all_terms_query,
        )
    ]

    shortened_terms = [
        quote_fts_token(
            token,
            shorten=True,
        )
        for token in tokens
    ]

    shortened_query = " AND ".join(shortened_terms)

    if shortened_query != all_terms_query:
        queries.append(
            (
                "prefix_terms",
                shortened_query,
            )
        )

    if len(full_terms) > 1:
        queries.append(
            (
                "any_term",
                " OR ".join(full_terms),
            )
        )

    return queries


def verify_search_schema(
    conn: sqlite3.Connection,
) -> None:
    fts_row = conn.execute("""
        SELECT 1
        FROM sqlite_master
        WHERE type = 'table'
          AND name = 'chunks_fts'
        """).fetchone()

    if fts_row is None:
        raise RuntimeError(
            "FTS5 index v databáze chýba. " "Spusti python -m tools.rag.rebuild_index."
        )

    embedding_row = conn.execute("""
        SELECT 1
        FROM sqlite_master
        WHERE type = 'table'
          AND name = 'chunk_embeddings'
        """).fetchone()

    if embedding_row is None:
        raise RuntimeError(
            "Embedding index v databáze chýba. " "Spusti python -m tools.rag.rebuild_index."
        )


def embedding_index_info(
    conn: sqlite3.Connection,
) -> tuple[str, int]:
    row = conn.execute("""
        SELECT
            model,
            dimensions
        FROM chunk_embeddings
        LIMIT 1
        """).fetchone()

    if row is None:
        raise RuntimeError("Embedding index je prázdny")
    if conn.execute("SELECT 1 FROM sqlite_master WHERE name='embedding_model_info'").fetchone():
        metadata = conn.execute("SELECT metadata_json FROM embedding_model_info").fetchone()
        if (
            metadata
            and json.loads(metadata[0])["fingerprint"]
            != model_profile(str(row["model"]))["fingerprint"]
        ):
            raise RuntimeError("Verzia embeddingového modelu nezodpovedá indexu; obnovte index.")

    return (
        str(row["model"]),
        int(row["dimensions"]),
    )


def make_source_url(
    document_path: str,
) -> str:
    clean_path = document_path

    if clean_path.startswith("pages/"):
        clean_path = clean_path[len("pages/") :]

    if clean_path.endswith("/README.md"):
        clean_path = clean_path[: -len("/README.md")]

    return "https://zp.kemt.fei.tuke.sk/" f"{clean_path}"


def load_labels(
    conn: sqlite3.Connection,
    table: str,
    column: str,
    chunk_ids: list[str],
) -> dict[str, list[str]]:
    if not chunk_ids:
        return {}

    placeholders = ",".join("?" for _ in chunk_ids)

    rows = conn.execute(
        f"""
        SELECT
            chunk_id,
            {column}
        FROM {table}
        WHERE chunk_id IN (
            {placeholders}
        )
        ORDER BY
            chunk_id,
            {column}
        """,
        chunk_ids,
    ).fetchall()

    values: dict[
        str,
        list[str],
    ] = defaultdict(list)

    for chunk_id, value in rows:
        values[chunk_id].append(value)

    return dict(values)


def run_fts_query(
    conn: sqlite3.Connection,
    match_query: str,
    candidate_limit: int,
    published_only: bool,
    *,
    eligible_chunk_ids: set[str] | None = None,
) -> list[dict[str, Any]]:
    rows = conn.execute(
        f"""
        SELECT
            chunks.chunk_id,
            chunks.document_path,
            chunks.title,
            chunks.author,
            chunks.published,
            chunks.chunk_index,
            chunks.heading_paths_json,
            chunks.text,
            chunks.text_length,
            chunks.token_count,
            chunks.content_hash,
            chunks_fts.rank AS bm25_score,
            snippet(
                chunks_fts,
                6,
                '',
                '',
                ' … ',
                36
            ) AS snippet
        FROM chunks_fts
        JOIN chunks
            ON chunks.id = chunks_fts.rowid
        WHERE chunks_fts MATCH ?
          AND chunks_fts.rank MATCH
              'bm25({BM25_SQL})'
          AND (
              ? = 0
              OR COALESCE(chunks.published, 1) = 1
          )
          AND (? IS NULL OR chunks.chunk_id IN (SELECT value FROM json_each(?)))
        ORDER BY
            chunks_fts.rank ASC,
            chunks.id ASC
        LIMIT ?
        """,
        (
            match_query,
            1 if published_only else 0,
            None if eligible_chunk_ids is None else json.dumps(sorted(eligible_chunk_ids)),
            None if eligible_chunk_ids is None else json.dumps(sorted(eligible_chunk_ids)),
            candidate_limit,
        ),
    ).fetchall()

    return [dict(row) for row in rows]


def run_vector_query(
    conn: sqlite3.Connection,
    query: str,
    candidate_limit: int,
    published_only: bool,
    *,
    eligible_chunk_ids: set[str] | None = None,
) -> list[dict[str, Any]]:
    if eligible_chunk_ids is not None and not eligible_chunk_ids:
        return []
    model_name, dimensions = embedding_index_info(conn)

    query_vector = embed_query(
        query,
        model_name=model_name,
    )

    if int(query_vector.shape[0]) != dimensions:
        raise RuntimeError("Rozmer query embeddingu " "sa nezhoduje s indexom")

    rows = conn.execute(
        """
        SELECT
            chunks.chunk_id,
            chunks.document_path,
            chunks.title,
            chunks.author,
            chunks.published,
            chunks.chunk_index,
            chunks.heading_paths_json,
            chunks.text,
            chunks.text_length,
            chunks.token_count,
            chunks.content_hash,
            chunk_embeddings.embedding,
            chunk_embeddings.dimensions
        FROM chunk_embeddings
        JOIN chunks
            ON chunks.chunk_id
            = chunk_embeddings.chunk_id
        WHERE (
            ? = 0
            OR COALESCE(chunks.published, 1) = 1
        )
        AND (? IS NULL OR chunks.chunk_id IN (SELECT value FROM json_each(?)))
        """,
        (1 if published_only else 0,
         None if eligible_chunk_ids is None else json.dumps(sorted(eligible_chunk_ids)),
         None if eligible_chunk_ids is None else json.dumps(sorted(eligible_chunk_ids))),
    ).fetchall()

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

    for row in rows:
        item = dict(row)

        blob = item.pop("embedding")

        stored_dimensions = int(item.pop("dimensions"))

        vector = blob_to_vector(
            blob,
            stored_dimensions,
        )

        similarity = cosine_similarity(
            query_vector,
            vector,
        )

        item["vector_score"] = round(
            similarity,
            6,
        )

        candidates.append(item)

    candidates.sort(
        key=lambda item: (
            -item["vector_score"],
            item["document_path"],
            item["chunk_index"],
        )
    )

    # A chunk can have several lossless embedding segments. Rank by its best
    # matching segment, and never spend several result slots on the same chunk.
    unique = {}
    for item in candidates:
        unique.setdefault(item["chunk_id"], item)
    return list(unique.values())[:candidate_limit]


def tokens_match_for_title(
    title_token: str,
    query_token: str,
) -> bool:
    if not title_token or not query_token:
        return False

    if title_token == query_token:
        return True

    shorter_length = min(
        len(title_token),
        len(query_token),
    )

    # Pri veľmi krátkych tokenoch nechceme
    # agresívne prefixové zhody.
    if shorter_length < 3:
        return False

    return title_token.startswith(query_token) or query_token.startswith(title_token)


def query_contains_title(
    query: str,
    title: str,
) -> bool:
    normalized_query = normalize_for_compare(query)

    normalized_title = normalize_for_compare(title)

    if not normalized_query or not normalized_title:
        return False

    title_tokens = normalized_title.split()

    query_token_values = normalized_query.split()

    # Jednoslovné názvy by mohli vytvárať
    # príliš veľa falošných boostov.
    if len(title_tokens) < 2:
        return False

    return all(
        any(
            tokens_match_for_title(
                title_token,
                query_token,
            )
            for query_token in query_token_values
        )
        for title_token in title_tokens
    )


def exact_match_bonus(
    query: str,
    item: dict[str, Any],
    tags: list[str],
    categories: list[str],
) -> float:
    normalized_query = normalize_for_compare(query)

    if not normalized_query:
        return 0.0

    raw_title = item.get("title") or ""

    title = normalize_for_compare(raw_title)

    author = normalize_for_compare(item.get("author") or "")

    path = normalize_for_compare(item.get("document_path") or "")

    text = normalize_for_compare(item.get("text") or "")

    normalized_tags = [normalize_for_compare(value) for value in tags]

    normalized_categories = [normalize_for_compare(value) for value in categories]

    bonus = 0.0

    # Pri otázkach typu:
    #
    # Kto sa venuje práci s názvom
    # „Systém získavania informácií
    # v slovenskom jazyku“?
    #
    # je celý názov práce veľmi silný
    # lexikálny signál, aj keď title
    # dokumentu je meno študenta.
    for phrase in quoted_query_phrases(query):
        if phrase in title or phrase in text:
            bonus += 6.0
            break

    if title == normalized_query:
        bonus += 6.0

    elif query_contains_title(
        query,
        raw_title,
    ):
        bonus += 5.0

    elif normalized_query in title:
        bonus += 3.0

    if author == normalized_query:
        bonus += 5.0

    elif normalized_query in author:
        bonus += 2.0

    if normalized_query in path:
        bonus += 2.0

    if normalized_query in (normalized_tags):
        bonus += 4.0

    if normalized_query in (normalized_categories):
        bonus += 3.0

    if normalized_query in text:
        bonus += 1.5

    return bonus


def database_bool(
    value: Any,
) -> bool | None:
    if value is None:
        return None

    return bool(value)


def parse_heading_paths(
    item: dict[str, Any],
) -> list:
    try:
        return json.loads(item.pop("heading_paths_json") or "[]")

    except json.JSONDecodeError:
        return []


def load_result_labels(
    conn: sqlite3.Connection,
    candidates: list[dict[str, Any]],
) -> tuple[
    dict[str, list[str]],
    dict[str, list[str]],
]:
    chunk_ids = [item["chunk_id"] for item in candidates]

    tags = load_labels(
        conn,
        "chunk_tags",
        "tag",
        chunk_ids,
    )

    categories = load_labels(
        conn,
        "chunk_categories",
        "category",
        chunk_ids,
    )

    return (
        tags,
        categories,
    )


def add_fts_metadata(
    conn: sqlite3.Connection,
    query: str,
    candidates: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    (
        tags_by_chunk,
        categories_by_chunk,
    ) = load_result_labels(
        conn,
        candidates,
    )

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

    for item in candidates:
        chunk_id = item["chunk_id"]

        tags = tags_by_chunk.get(
            chunk_id,
            [],
        )

        categories = categories_by_chunk.get(
            chunk_id,
            [],
        )

        bm25_score = float(item.pop("bm25_score"))

        strategy = item.pop("strategy")

        base_score = max(
            0.0,
            -bm25_score,
        )

        match_bonus = exact_match_bonus(
            query,
            item,
            tags,
            categories,
        )

        score = base_score + match_bonus

        heading_paths = parse_heading_paths(item)

        item["published"] = database_bool(item.get("published"))

        item["_strategy_priority"] = STRATEGY_PRIORITY[strategy]

        item.update(
            {
                "heading_paths": (heading_paths),
                "tags": tags,
                "categories": (categories),
                "score": round(
                    score,
                    6,
                ),
                "bm25_score": round(
                    bm25_score,
                    6,
                ),
                "match_strategy": (strategy),
                "source_url": (make_source_url(item["document_path"])),
            }
        )

        results.append(item)

    results.sort(
        key=lambda item: (
            -item["_strategy_priority"],
            -item["score"],
            item["bm25_score"],
            item["document_path"],
            item["chunk_index"],
        )
    )

    for item in results:
        item.pop(
            "_strategy_priority",
            None,
        )

    return results


def add_vector_metadata(
    conn: sqlite3.Connection,
    candidates: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    (
        tags_by_chunk,
        categories_by_chunk,
    ) = load_result_labels(
        conn,
        candidates,
    )

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

    for item in candidates:
        chunk_id = item["chunk_id"]

        item["published"] = database_bool(item.get("published"))

        item["heading_paths"] = parse_heading_paths(item)

        item["tags"] = tags_by_chunk.get(
            chunk_id,
            [],
        )

        item["categories"] = categories_by_chunk.get(
            chunk_id,
            [],
        )

        item["source_url"] = make_source_url(item["document_path"])

        text = item.get("text") or ""

        item["snippet"] = text[:320].strip()

        results.append(item)

    return results


def lexical_bonus_from_fts_result(
    item: dict[str, Any],
) -> float:
    fts_score = item.get("score")

    bm25_score = item.get("bm25_score")

    if fts_score is None or bm25_score is None:
        return 0.0

    try:
        score_value = float(fts_score)

        bm25_value = float(bm25_score)

    except (
        TypeError,
        ValueError,
    ):
        return 0.0

    base_score = max(
        0.0,
        -bm25_value,
    )

    return max(
        0.0,
        score_value - base_score,
    )


def fuse_hybrid_results(
    fts_results: list[dict[str, Any]],
    vector_results: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    merged: dict[
        str,
        dict[str, Any],
    ] = {}

    scores: dict[
        str,
        float,
    ] = defaultdict(float)

    fts_ranks: dict[
        str,
        int,
    ] = {}

    vector_ranks: dict[
        str,
        int,
    ] = {}

    for rank, item in enumerate(
        fts_results,
        start=1,
    ):
        chunk_id = item["chunk_id"]

        merged[chunk_id] = dict(item)

        fts_ranks[chunk_id] = rank

        strategy = item.get("match_strategy")

        fts_weight = ANY_TERM_RRF_WEIGHT if strategy == "any_term" else FTS_RRF_WEIGHT

        lexical_bonus = lexical_bonus_from_fts_result(item)

        # Silný lexikálny anchor má prednosť
        # pred čisto sémantickou podobnosťou.
        #
        # Typické prípady:
        # - meno osoby je priamo title dokumentu,
        # - názov práce je uvedený v úvodzovkách,
        # - query obsahuje celý názov dokumentu.
        #
        # Toto rieši najmä situáciu, keď
        # any_term FTS nájde správneho Bogdana
        # na 1. mieste, ale pôvodná váha 0.25
        # ho po RRF vytlačila mimo Top 5.
        if lexical_bonus >= LEXICAL_ANCHOR_MIN_BONUS:
            fts_weight = max(
                fts_weight,
                LEXICAL_ANCHOR_RRF_WEIGHT,
            )

        scores[chunk_id] += fts_weight / (RRF_K + rank)

    for rank, item in enumerate(
        vector_results,
        start=1,
    ):
        chunk_id = item["chunk_id"]

        vector_ranks[chunk_id] = rank

        scores[chunk_id] += VECTOR_RRF_WEIGHT / (RRF_K + rank)

        if chunk_id not in merged:
            merged[chunk_id] = dict(item)

        else:
            merged[chunk_id]["vector_score"] = item["vector_score"]

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

    for chunk_id, item in merged.items():
        item["fts_score"] = item.get("score")

        item["fts_rank"] = fts_ranks.get(chunk_id)

        item["vector_rank"] = vector_ranks.get(chunk_id)

        item.setdefault(
            "vector_score",
            None,
        )

        item.setdefault(
            "bm25_score",
            None,
        )

        item.setdefault(
            "match_strategy",
            None,
        )

        hybrid_score = scores[chunk_id]

        item["hybrid_score"] = round(
            hybrid_score,
            8,
        )

        item["score"] = round(
            hybrid_score,
            8,
        )

        results.append(item)

    results.sort(
        key=lambda item: (
            -item["hybrid_score"],
            item["document_path"],
            item["chunk_index"],
        )
    )

    return results


def diversify_results(
    results: list[dict[str, Any]],
    limit: int,
    max_per_document: int,
) -> list[dict[str, Any]]:
    if max_per_document <= 0:
        return results[:limit]

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

    document_counts: dict[
        str,
        int,
    ] = defaultdict(int)

    for item in results:
        document_path = item["document_path"]

        if document_counts[document_path] >= max_per_document:
            continue

        selected.append(item)

        document_counts[document_path] += 1

        if len(selected) >= limit:
            break

    return selected