from __future__ import annotations

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

from scripts.embedding_utils import (
    blob_to_vector,
    cosine_similarity,
    embed_query,
)


WORD_RE = re.compile(
    r"[^\W_]+",
    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


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 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 scripts/rebuild_index.py."
        )

    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 scripts/rebuild_index.py."
        )


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"
        )

    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,
) -> 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 chunks.published = 1
          )
        ORDER BY
            chunks_fts.rank ASC,
            chunks.id ASC
        LIMIT ?
        """,
        (
            match_query,
            1 if published_only else 0,
            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,
) -> list[dict[str, Any]]:
    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 chunks.published = 1
        )
        """,
        (
            1 if published_only else 0,
        ),
    ).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"],
        )
    )

    return candidates[
        :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

    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,
        )

        score = (
            base_score
            + exact_match_bonus(
                query,
                item,
                tags,
                categories,
            )
        )

        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 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
        )

        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