"""
Search algorithms for Lattice.

Core layer - pure logic for hybrid search (BM25 + Vector + RRF).
No I/O operations.

Reference: RFC-002 §5.2 Hybrid Search (RRF)
"""

from __future__ import annotations

import deal

from lattice.core.types.search import FtsResult, RankedItem, VecResult

# RRF constant (standard value from research)
RRF_K = 60


# =============================================================================
# Positional Ranking
# =============================================================================


@deal.post(lambda result: isinstance(result, dict))
@deal.post(lambda result: all(isinstance(k, int) for k in result.keys()))
@deal.post(lambda result: all(isinstance(v, int) for v in result.values()))
def assign_positional_ranks(
    results: list[FtsResult] | list[VecResult],
) -> dict[int, int]:
    """Assign positional ranks (1-based) to results based on sort order.

    For FTS results: lower rank (more negative) = higher position.
    For Vec results: lower distance = higher position.

    Args:
        results: List of FTS or Vector search results (already sorted).

    Returns:
        Dictionary mapping rowid to positional rank (1-based).

    Example:
        >>> from lattice.core.types.search import FtsResult
        >>> results = [FtsResult(rowid=1, rank=-0.5), FtsResult(rowid=2, rank=-0.3)]
        >>> ranks = assign_positional_ranks(results)
        >>> ranks[1]
        1
        >>> ranks[2]
        2
    """
    ranks: dict[int, int] = {}
    for position, result in enumerate(results, start=1):
        ranks[result.rowid] = position
    return ranks


# =============================================================================
# RRF Score Calculation
# =============================================================================


@deal.pre(lambda bm25_rank, vec_rank, k=RRF_K: k >= 1)
@deal.pre(lambda bm25_rank, vec_rank, k=RRF_K: bm25_rank is None or bm25_rank >= 1)
@deal.pre(lambda bm25_rank, vec_rank, k=RRF_K: vec_rank is None or vec_rank >= 1)
@deal.post(lambda result: result >= 0)
def calculate_rrf_score(
    bm25_rank: int | None,
    vec_rank: int | None,
    k: int = RRF_K,
) -> float:
    """Calculate RRF (Reciprocal Rank Fusion) score.

    RRF Score = 1/(k + bm25_rank) + 1/(k + vec_rank)

    A result appearing in both lists gets a higher score.
    A result appearing in only one list gets partial credit.

    Args:
        bm25_rank: Positional rank from BM25 (1-based), None if not in BM25 results.
        vec_rank: Positional rank from vector search (1-based), None if not in vector results.
        k: RRF constant (default 60, standard from research). Must be >= 1.

    Returns:
        RRF score (higher is better).

    Example:
        >>> score = calculate_rrf_score(1, 1)  # Top result in both
        >>> round(score, 4)
        0.0328
        >>> score = calculate_rrf_score(1, None)  # Only in BM25
        >>> round(score, 4)
        0.0164
        >>> calculate_rrf_score(None, None)  # In neither
        0.0
    """
    score = 0.0
    if bm25_rank is not None:
        score += 1.0 / (k + bm25_rank)
    if vec_rank is not None:
        score += 1.0 / (k + vec_rank)
    return score


# =============================================================================
# Result Fusion
# =============================================================================


@deal.post(lambda result: isinstance(result, list))
@deal.post(lambda result: all(isinstance(item, RankedItem) for item in result))
@deal.post(lambda result: len(result) == len(set(item.rowid for item in result)))
def fuse_results(
    bm25_results: list[FtsResult],
    vec_results: list[VecResult],
    k: int = RRF_K,
) -> list[RankedItem]:
    """Fuse BM25 and vector search results using RRF.

    Steps:
    1. Assign positional ranks to each result set.
    2. Calculate RRF score for each unique rowid.
    3. Sort by RRF score (descending).

    Args:
        bm25_results: Results from FTS5 full-text search.
        vec_results: Results from sqlite-vec vector search.
        k: RRF constant (default 60).

    Returns:
        List of RankedItem sorted by RRF score (descending).

    Example:
        >>> from lattice.core.types.search import FtsResult, VecResult
        >>> bm25 = [FtsResult(rowid=1, rank=-0.5), FtsResult(rowid=2, rank=-0.3)]
        >>> vec = [VecResult(rowid=1, distance=0.1), VecResult(rowid=3, distance=0.2)]
        >>> fused = fuse_results(bm25, vec)
        >>> fused[0].rowid  # rowid=1 appears in both, highest score
        1
    """
    # Step 1: Assign positional ranks
    bm25_ranks = assign_positional_ranks(bm25_results)
    vec_ranks = assign_positional_ranks(vec_results)

    # Step 2: Get all unique rowids
    all_rowids = get_unique_rowids(bm25_results, vec_results)

    # Step 3: Calculate RRF scores
    scored: list[RankedItem] = []
    for rowid in all_rowids:
        bm25_rank = bm25_ranks.get(rowid)
        vec_rank = vec_ranks.get(rowid)
        rrf_score = calculate_rrf_score(bm25_rank, vec_rank, k)
        scored.append(RankedItem(rowid=rowid, rrf_score=rrf_score))

    # Step 4: Sort by RRF score descending
    scored.sort(key=lambda x: x.rrf_score, reverse=True)

    return scored


@deal.post(lambda result: isinstance(result, list))
@deal.post(lambda result: all(isinstance(item, RankedItem) for item in result))
def top_n(results: list[RankedItem], n: int) -> list[RankedItem]:
    """Return top N results from a sorted list.

    Results should already be sorted by RRF score (descending).

    Args:
        results: List of RankedItem (assumed sorted by score descending).
        n: Number of results to return. If n > len(results), returns all.

    Returns:
        Top N results.

    Example:
        >>> items = [RankedItem(1, 0.9), RankedItem(2, 0.5), RankedItem(3, 0.3)]
        >>> top_n(items, 2)
        [RankedItem(rowid=1, rrf_score=0.9), RankedItem(rowid=2, rrf_score=0.5)]
    """
    if n <= 0:
        return []
    return results[:n]


# =============================================================================
# Utility Functions
# =============================================================================


@deal.post(lambda result: isinstance(result, list))
@deal.post(lambda result: len(result) == len(set(result)))  # No duplicates
def get_unique_rowids(
    bm25_results: list[FtsResult],
    vec_results: list[VecResult],
) -> list[int]:
    """Get all unique rowids from both result sets.

    Args:
        bm25_results: Results from FTS5 search.
        vec_results: Results from vector search.

    Returns:
        List of unique rowids (order preserved, BM25 first then vector).

    Example:
        >>> from lattice.core.types.search import FtsResult, VecResult
        >>> bm25 = [FtsResult(rowid=1, rank=-0.5), FtsResult(rowid=2, rank=-0.3)]
        >>> vec = [VecResult(rowid=1, distance=0.1), VecResult(rowid=3, distance=0.2)]
        >>> get_unique_rowids(bm25, vec)
        [1, 2, 3]
    """
    seen: set[int] = set()
    unique: list[int] = []

    # Add BM25 rowids first
    for result in bm25_results:
        if result.rowid not in seen:
            seen.add(result.rowid)
            unique.append(result.rowid)

    # Add vector rowids
    for result in vec_results:
        if result.rowid not in seen:
            seen.add(result.rowid)
            unique.append(result.rowid)

    return unique
