"""
Reciprocal Rank Fusion (RRF) utilities for combining search results.

RRF is a method to combine multiple ranked lists by computing the reciprocal
of each item's rank in each list, then summing these reciprocal ranks.
"""

import re
from typing import Any, TypeVar

from sqlalchemy import Select, and_, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from src import models
from src.config import settings
from src.embedding_client import embedding_client
from src.exceptions import ValidationException
from src.models import session_peers_table
from src.utils.filter import apply_filter
from src.utils.formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern
from src.vector_store import get_external_vector_store

T = TypeVar("T")


def reciprocal_rank_fusion(*ranked_lists: list[T], k: int = 60, limit: int) -> list[T]:
    """
    Combine multiple ranked lists using Reciprocal Rank Fusion (RRF).

    RRF assigns a score to each item based on the formula:
    RRF_score = sum(1 / (k + rank_i)) for all lists where the item appears

    Where:
    - k is a constant (typically 60) that controls the impact of high-ranked items
    - rank_i is the rank of the item in list i (1-indexed)

    Args:
        *ranked_lists: Variable number of ranked lists to combine
        k: RRF constant parameter (default: 60)
        limit: Maximum number of results to return

    Returns:
        list of items ranked by RRF score (highest score first)
    """
    if not ranked_lists:
        return []

    # dictionary to store RRF scores for each item
    rrf_scores: dict[T, float] = {}

    # Process each ranked list
    for ranked_list in ranked_lists:
        for rank, item in enumerate(ranked_list, 1):  # 1-indexed ranking
            if item not in rrf_scores:
                rrf_scores[item] = 0.0
            # Add reciprocal rank contribution from this list
            rrf_scores[item] += 1.0 / (k + rank)

    # Sort items by RRF score (descending order)
    sorted_items = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)

    # Extract just the items (not the scores)
    result = [item for item, _ in sorted_items]

    return result[:limit]


async def _semantic_search(
    db: AsyncSession,
    query: str,
    workspace_name: str,
    limit: int,
    filters: dict[str, Any] | None = None,
) -> list[models.Message]:
    """
    Perform semantic search using external vector store for message embeddings.

    Args:
        db: Database session
        query: Search query
        workspace_name: Name of the workspace to search in
        limit: Maximum number of results to return
        filters: Optional filters to apply at vector store level (supports: session_id, peer_id)

    Returns:
        list of messages ordered by semantic similarity
    """
    try:
        embedding_query = await embedding_client.embed(query)
    except ValueError as e:
        raise ValidationException(
            f"Query exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}."
        ) from e

    # Query Postgres / pgvector directly
    if settings.EMBED_MESSAGES and (
        settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED
    ):
        # Join message_embeddings with messages to get full message objects
        distance_expr = models.MessageEmbedding.embedding.cosine_distance(
            embedding_query
        )

        stmt = (
            select(models.Message)
            .join(
                models.MessageEmbedding,
                models.Message.public_id == models.MessageEmbedding.message_id,
            )
            .where(models.MessageEmbedding.embedding.isnot(None))
            .where(models.MessageEmbedding.workspace_name == workspace_name)
        )

        # Apply all additional filters using the standard filter utility
        # filters dict uses external names (session_id, peer_id) which apply_filter will map
        # to internal column names (session_name, peer_name)
        if filters:
            # Create a copy with workspace added
            internal_filters = filters.copy()
            internal_filters["workspace_id"] = workspace_name
            stmt = apply_filter(stmt, models.Message, internal_filters)

        # Order by cosine distance and limit
        stmt = stmt.order_by(distance_expr).limit(limit)

        result = await db.execute(stmt)
        return list(result.scalars().all())

    # FALLBACK: Use external vector store (Turbopuffer, LanceDB)
    external_vector_store = get_external_vector_store()
    if external_vector_store is None:
        return []

    namespace = external_vector_store.get_vector_namespace("message", workspace_name)

    # Build vector store filters from the provided filters
    vector_filters: dict[str, Any] = {}
    if filters:
        # Map external filter keys to vector store metadata keys
        if "session_id" in filters:
            vector_filters["session_name"] = filters["session_id"]
        if "peer_id" in filters:
            vector_filters["peer_name"] = filters["peer_id"]

    # Query external vector store for similar message embeddings
    # Since all filters are applied at the vector store level, we don't need to oversample
    vector_results = await external_vector_store.query(
        namespace,
        embedding_query,
        top_k=limit,
        filters=vector_filters if vector_filters else None,
    )

    if not vector_results:
        return []

    # Extract message IDs from vector metadata
    # Use dict to deduplicate while preserving order (dict keys maintain insertion order in Python 3.7+)
    seen_message_ids: dict[str, None] = {}

    for result in vector_results:
        message_id = result.metadata.get("message_id")
        if message_id and message_id not in seen_message_ids:
            seen_message_ids[message_id] = None

    message_ids = list(seen_message_ids.keys())

    # Fetch messages from database by the IDs from vector search and reapply filters
    semantic_query = select(models.Message).where(
        models.Message.public_id.in_(message_ids)
    )
    semantic_query = apply_filter(semantic_query, models.Message, filters)

    result = await db.execute(semantic_query)
    messages = {msg.public_id: msg for msg in result.scalars().all()}

    # Return messages in order of similarity (preserving vector store order)
    ordered_messages: list[models.Message] = []
    for msg_id in message_ids:
        if msg_id in messages:
            ordered_messages.append(messages[msg_id])

    return ordered_messages


async def _filter_by_peer_perspective(
    db: AsyncSession,
    messages: list[models.Message],
    workspace_name: str,
    peer_name: str,
) -> list[models.Message]:
    """
    Filter messages by peer perspective (temporal session membership).

    Only keeps messages from sessions where the peer was a member at the time
    the message was created (between joined_at and left_at).

    Args:
        db: Database session
        messages: List of messages to filter
        workspace_name: Name of the workspace
        peer_name: Name of the peer whose perspective to use

    Returns:
        Filtered list of messages
    """
    if not messages:
        return []

    # Get all session memberships for this peer in this workspace
    session_memberships_query = (
        select(session_peers_table)
        .where(session_peers_table.c.workspace_name == workspace_name)
        .where(session_peers_table.c.peer_name == peer_name)
    )
    result = await db.execute(session_memberships_query)
    memberships = result.all()

    # Build a lookup of session -> time windows
    session_windows: dict[str, list[tuple[Any, Any]]] = {}
    for membership in memberships:
        session_name = membership.session_name
        if session_name not in session_windows:
            session_windows[session_name] = []
        session_windows[session_name].append((membership.joined_at, membership.left_at))

    # Filter messages
    filtered_messages: list[models.Message] = []
    for msg in messages:
        if msg.session_name not in session_windows:
            continue

        # Check if message was created during any of the peer's active windows in this session
        for joined_at, left_at in session_windows[msg.session_name]:
            if msg.created_at >= joined_at and (
                left_at is None or msg.created_at <= left_at
            ):
                filtered_messages.append(msg)
                break  # Don't add the same message twice

    return filtered_messages


async def _fulltext_search(
    db: AsyncSession,
    query: str,
    stmt: Select[tuple[models.Message]],
    limit: int,
) -> list[models.Message]:
    """
    Perform full-text search using PostgreSQL FTS and ILIKE fallback.

    Args:
        db: Database session
        query: Search query
        stmt: Base SQL query conditions
        limit: Maximum number of results to return

    Returns:
        list of messages ordered by text search relevance
    """
    # Check if query contains special characters that FTS might not handle well
    has_special_chars = bool(
        re.search(r'[~`!@#$%^&*()_+=\[\]{};\':"\\|,.<>/?-]', query)
    )

    # Escape ILIKE pattern characters to treat user input literally
    escaped_query = escape_ilike_pattern(query)

    if has_special_chars:
        # For queries with special characters, use exact string matching (ILIKE)
        search_condition = models.Message.content.ilike(
            f"%{escaped_query}%", escape=ILIKE_ESCAPE_CHAR
        )
        fulltext_query = stmt.where(search_condition).order_by(
            models.Message.created_at.desc()
        )
    else:
        # For natural language queries, use full text search with ranking
        fts_condition = func.to_tsvector("english", models.Message.content).op("@@")(
            func.plainto_tsquery("english", query)
        )

        # Combine FTS with ILIKE as fallback for better coverage
        combined_condition = or_(
            fts_condition,
            models.Message.content.ilike(
                f"%{escaped_query}%", escape=ILIKE_ESCAPE_CHAR
            ),
        )

        fulltext_query = stmt.where(combined_condition).order_by(
            # Order by FTS relevance first, then by creation time
            func.coalesce(
                func.ts_rank(
                    func.to_tsvector("english", models.Message.content),
                    func.plainto_tsquery("english", query),
                ),
                0,
            ).desc(),
            models.Message.created_at.desc(),
        )

    fulltext_query = fulltext_query.limit(limit)

    result = await db.execute(fulltext_query)
    return list(result.scalars().all())


async def search(
    db: AsyncSession,
    query: str,
    *,
    filters: dict[str, Any] | None = None,
    limit: int = 10,
) -> list[models.Message]:
    """
    Search across message content using a hybrid approach with Reciprocal Rank Fusion (RRF).

    This function combines semantic search and full-text search results using RRF when both
    are available, providing better search results than either method alone.

    Args:
        db: Database session
        query: Search query to match against message content
        filters: Optional filters to scope search (must include workspace_id for semantic search).
            Special filter 'peer_perspective' will search across all messages from sessions that the peer is/was a member of,
            filtered by the time window when they were actually in the session.
        limit: Maximum number of results to return

    Returns:
        list of messages that match the search query, ordered by RRF relevance or individual search relevance

    Raises:
        ValidationException: If query exceeds maximum token limit for embeddings
    """
    # Base query conditions
    stmt = select(models.Message)

    # Handle special peer_perspective filter
    peer_perspective_name: str | None = None
    if filters and "peer_perspective" in filters:
        peer_perspective_name = filters["peer_perspective"]
        # Remove from filters dict so apply_filter doesn't try to handle it
        filters = {k: v for k, v in filters.items() if k != "peer_perspective"}
        # Safety: peer_perspective must be scoped to a workspace
        if not filters or (
            "workspace_id" not in filters and "workspace_name" not in filters
        ):
            raise ValidationException(
                "peer_perspective requires a workspace scope (workspace_id or workspace_name)."
            )

        # Join with session_peers_table to get messages from sessions the peer was in
        # Only include messages created during the time window the peer was active
        stmt = stmt.join(
            session_peers_table,
            and_(
                models.Message.session_name == session_peers_table.c.session_name,
                models.Message.workspace_name == session_peers_table.c.workspace_name,
                models.Message.created_at >= session_peers_table.c.joined_at,
                or_(
                    session_peers_table.c.left_at.is_(None),
                    models.Message.created_at <= session_peers_table.c.left_at,
                ),
            ),
        ).where(session_peers_table.c.peer_name == peer_perspective_name)

    stmt = apply_filter(stmt, models.Message, filters)

    search_results: list[list[models.Message]] = []

    # Perform semantic search if enabled and we have workspace context
    # workspace_id is required for semantic search to determine the vector namespace
    workspace_name: str | None = filters.get("workspace_id") if filters else None
    if settings.EMBED_MESSAGES and isinstance(workspace_name, str):
        # Type narrowing: workspace_name is guaranteed to be str in this block
        # Get more results for fusion (increase if peer_perspective filtering is applied post-search)
        semantic_limit = limit * 4 if peer_perspective_name else limit * 2
        semantic_results = await _semantic_search(
            db=db,
            query=query,
            workspace_name=workspace_name,
            limit=semantic_limit,
            filters=filters,
        )

        # Apply peer_perspective filtering to semantic results if needed
        # Vector store can't handle temporal filtering (joined_at/left_at), so filter post-search
        if peer_perspective_name:
            semantic_results = await _filter_by_peer_perspective(
                db, semantic_results, workspace_name, peer_perspective_name
            )

        search_results.append(semantic_results)

    # Perform full-text search
    # Get more results for fusion
    fulltext_limit = limit * 2
    fulltext_results = await _fulltext_search(
        db=db, query=query, stmt=stmt, limit=fulltext_limit
    )
    search_results.append(fulltext_results)

    # Combine results using RRF if we have multiple search methods
    if len(search_results) > 1:
        # Use RRF to combine semantic and full-text results
        combined_results = reciprocal_rank_fusion(*search_results, limit=limit)
    elif len(search_results) == 1:
        # Single search method - apply limit directly
        combined_results = search_results[0]
        combined_results = combined_results[:limit]
    else:
        # No search results
        combined_results = []

    return combined_results