"""Hybrid search orchestration."""

from __future__ import annotations

import asyncio
import logging
import re
from collections import defaultdict
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING

from pgvector.sqlalchemy import Vector
from sqlalchemy import literal, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession

from python.ebook_search.bm25_corpus import (
    BM25CorpusUnavailableError,
    load_bm25_corpus,
    score_bm25_corpus,
)
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
from python.ebook_search.protected_phrases.matching import (
    detect_protected_phrases_for_query,
    phrase_hits_for_chunks,
)
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.timing import RuntimeStep, async_timed_result, timed_result
from python.orm.richie import (
    EbookChapter,
    EbookChunk,
    EbookEmbeddingModel,
    EbookSource,
)

if TYPE_CHECKING:
    from collections.abc import Mapping, Sequence

    import httpx
    from sqlalchemy.ext.asyncio import AsyncEngine

    from python.ebook_search.config import EbookSearchConfig
    from python.ebook_search.protected_phrases.models import HydratedPhraseMatch

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SearchResult:
    """One source chunk returned by search."""

    chunk_id: int
    text: str
    source_title: str
    source_id: int | None = None
    score: float = 0.0
    vector_score: float | None = None
    bm25_score: float | None = None
    fused_score: float | None = None
    rerank_score: float | None = None
    phrase_hit_count: int = 0
    matched_phrases: tuple[str, ...] = ()
    source_author: str | None = None
    chapter_title: str | None = None
    page_label: str | None = None
    rank_source: str = "Hybrid"


@dataclass(frozen=True)
class SearchResponse:
    """Search output for the UI."""

    query: str
    results: list[SearchResult]
    rank_label: str
    timings: tuple[RuntimeStep, ...] = ()
    phrase_matches: tuple[HydratedPhraseMatch, ...] = ()

    @property
    def total_runtime_ms(self) -> float:
        """Return total measured runtime for the response."""
        return sum(step.duration_ms for step in self.timings if step.counts_toward_total)


@dataclass(frozen=True)
class RetrievalResponse:
    """Parallel retrieval output for vector and BM25 candidates."""

    vector_results: list[SearchResult]
    lexical_results: list[SearchResult]
    timings: tuple[RuntimeStep, ...]


async def search_ebooks(
    engine: AsyncEngine,
    client: httpx.AsyncClient,
    query: str,
    config: EbookSearchConfig,
    *,
    rerank: bool = False,
    phrase_matching: bool | None = None,
) -> SearchResponse:
    """Run hybrid vector/BM25 search and optional reranking."""
    if not query.strip():
        logger.info("ebook_search_empty_query")
        return SearchResponse(query=query, results=[], rank_label="Hybrid")

    phrase_matching_enabled = config.phrase_matching_enabled if phrase_matching is None else phrase_matching
    logger.info(
        "ebook_search_start query_length=%s rerank=%s phrase_matching=%s",
        len(query),
        rerank,
        phrase_matching_enabled,
    )
    timings: list[RuntimeStep] = []
    if phrase_matching_enabled:
        phrase_matches, timing = await async_timed_result(
            "Protected phrase detection", query_phrase_matches(engine, query, config)
        )
    else:
        phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
    timings.append(timing)
    retrieval, timing = await async_timed_result(
        "Hybrid retrieval",
        parallel_retrieval(engine, client, query, config),
    )
    timings.extend(retrieval.timings)
    timings.append(timing)
    fused, timing = timed_result(
        "Reciprocal rank fusion",
        reciprocal_rank_fusion,
        retrieval.vector_results,
        retrieval.lexical_results,
        rank_constant=config.rrf_rank_constant,
    )
    timings.append(timing)
    if phrase_matching_enabled:
        fused, timing = await async_timed_result(
            "Phrase mention boost",
            apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
        )
    else:
        fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused)
    timings.append(timing)
    if config.rerank.enabled and rerank:
        response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
    else:
        response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
    timings.append(timing)
    response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches))
    logger.info(
        "ebook_search_complete vector_candidates=%s lexical_candidates=%s "
        "fused_candidates=%s phrase_matching=%s phrase_matches=%s returned=%s rank_label=%s runtime_ms=%.1f",
        len(retrieval.vector_results),
        len(retrieval.lexical_results),
        len(fused),
        phrase_matching_enabled,
        len(phrase_matches),
        len(response.results),
        response.rank_label,
        response.total_runtime_ms,
    )
    return response


def skip_phrase_matches() -> list[HydratedPhraseMatch]:
    """Return no protected phrase matches when phrase matching is disabled."""
    logger.info("ebook_protected_phrase_detection_skipped")
    return []


async def query_phrase_matches(
    engine: AsyncEngine,
    query: str,
    config: EbookSearchConfig,
) -> list[HydratedPhraseMatch]:
    """Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
    try:
        async with AsyncSession(engine) as session:
            return await detect_protected_phrases_for_query(session, query, config)
    except SQLAlchemyError as error:
        logger.warning("ebook_protected_phrase_detection_unavailable error=%s", error)
        return []


def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchResult]:
    """Return candidates unchanged when phrase matching is disabled."""
    logger.info("ebook_phrase_boost_skipped candidates=%s", len(candidates))
    return candidates


async def apply_phrase_mention_boosts(
    engine: AsyncEngine,
    candidates: list[SearchResult],
    phrase_matches: Sequence[HydratedPhraseMatch],
    phrase_hit_boost: float,
) -> list[SearchResult]:
    """Boost retrieved chunks that have indexed mentions for detected protected phrases."""
    phrase_ids = sorted({match.phrase_id for match in phrase_matches})
    if not candidates or not phrase_ids or phrase_hit_boost <= 0:
        return candidates

    chunk_ids = [candidate.chunk_id for candidate in candidates]
    try:
        async with AsyncSession(engine) as session:
            phrase_hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
    except SQLAlchemyError as error:
        logger.warning("ebook_phrase_boost_unavailable error=%s", error)
        return candidates

    if not phrase_hits:
        return candidates

    hit_counts = {
        chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in phrase_hits.items()
    }
    boosted = [
        replace(
            candidate,
            score=candidate.score + (hit_counts.get(candidate.chunk_id, 0) * phrase_hit_boost),
            fused_score=boosted_fused_score(candidate, hit_counts.get(candidate.chunk_id, 0), phrase_hit_boost),
            phrase_hit_count=hit_counts.get(candidate.chunk_id, 0),
            matched_phrases=tuple(hit.phrase_text for hit in phrase_hits.get(candidate.chunk_id, ())),
            rank_source=phrase_rank_source(candidate.rank_source, hit_counts.get(candidate.chunk_id, 0)),
        )
        for candidate in candidates
    ]
    return sorted(boosted, key=lambda candidate: candidate.score, reverse=True)


def boosted_fused_score(candidate: SearchResult, phrase_hit_count: int, phrase_hit_boost: float) -> float | None:
    """Return a fused score adjusted by phrase hits when a fused score exists."""
    if candidate.fused_score is None:
        return None
    return candidate.fused_score + (phrase_hit_count * phrase_hit_boost)


def phrase_rank_source(rank_source: str, phrase_hit_count: int) -> str:
    """Append phrase evidence to a rank-source label when a chunk was boosted."""
    if phrase_hit_count <= 0 or "phrases" in rank_source:
        return rank_source
    return f"{rank_source} + phrases"


async def parallel_retrieval(
    engine: AsyncEngine,
    client: httpx.AsyncClient,
    query: str,
    config: EbookSearchConfig,
) -> RetrievalResponse:
    """Run vector and BM25 candidate retrieval concurrently with separate database sessions.

    BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
    instead of on the event loop.
    """
    (vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
        async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
        async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
    )

    logger.info(
        "ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
        len(vector_results),
        len(lexical_results),
    )
    return RetrievalResponse(
        vector_results=vector_results,
        lexical_results=lexical_results,
        timings=(
            replace(vector_timing, counts_toward_total=False),
            replace(lexical_timing, counts_toward_total=False),
        ),
    )


def skip_rerank(
    query: str,
    candidates: list[SearchResult],
    config: EbookSearchConfig,
) -> SearchResponse:
    """Return fused hybrid results without reranking."""
    logger.info("ebook_rerank_skipped candidates=%s", len(candidates))
    return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")


async def apply_rerank(
    client: httpx.AsyncClient,
    query: str,
    candidates: list[SearchResult],
    config: EbookSearchConfig,
) -> SearchResponse:
    """Rerank already-fused hybrid candidates."""
    reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
    logger.info(
        "ebook_rerank_complete input_candidates=%s returned=%s",
        min(len(candidates), config.rerank.candidates),
        len(reranked),
    )
    return SearchResponse(
        query=query,
        results=[replace(result, rank_source="Hybrid + rerank") for result in reranked[: config.top_k]],
        rank_label="Hybrid + rerank",
    )


async def vector_candidates(
    engine: AsyncEngine,
    client: httpx.AsyncClient,
    query: str,
    config: EbookSearchConfig,
) -> list[SearchResult]:
    """Return pgvector cosine candidates for a natural-language query."""
    async with AsyncSession(engine) as session:
        model = await session.scalar(
            select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)
        )
        if model is None:
            msg = f"Embedding model is not registered: {config.embedding_model}"
            raise ValueError(msg)

        expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
        if model.dimension != expected_dimension:
            msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
            raise ValueError(msg)

        embedding = await embed_query(client, query, config)
        limit = max(config.rerank.candidates, config.top_k) * config.vector_candidate_multiplier
        embedding_table = get_embedding_table(model.dimension)

        embedding_param = literal(embedding, type_=Vector(model.dimension))
        distance = embedding_table.embedding.op("<=>")(embedding_param)
        score = (literal(1.0) - distance).label("score")
        statement = (
            select(
                EbookChunk.id.label("chunk_id"),
                EbookChunk.text.label("text"),
                EbookSource.id.label("source_id"),
                EbookSource.title.label("source_title"),
                EbookSource.author.label("source_author"),
                EbookChapter.title.label("chapter_title"),
                EbookChunk.page_label.label("page_label"),
                score,
            )
            .select_from(embedding_table)
            .join(EbookChunk, EbookChunk.id == embedding_table.chunk_id)
            .join(EbookSource, EbookSource.id == EbookChunk.source_id)
            .outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id)
            .where(embedding_table.model_id == model.id)
            .order_by(distance)
            .limit(limit)
        )
        rows = (await session.execute(statement)).mappings()
        results = [search_result_from_row(row) for row in rows]
        logger.info(
            "ebook_vector_search_complete model=%s dimension=%s candidates=%s",
            config.embedding_model,
            model.dimension,
            len(results),
        )
        return results


def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]:
    """Return BM25-ranked lexical candidates using the persisted corpus."""
    try:
        corpus = load_bm25_corpus(config)
    except BM25CorpusUnavailableError as error:
        logger.warning("ebook_bm25_index_unavailable_skipping error=%s", error)
        return []

    if not corpus.records:
        logger.info("ebook_bm25_search_complete corpus=0 candidates=0")
        return []

    bm25_query = retrieval_query_from_text(query)
    scored_records = score_bm25_corpus(bm25_query, corpus, limit=config.bm25_candidate_limit)
    results = [
        replace(search_result_from_row(record), score=score, vector_score=None, bm25_score=score)
        for record, score in scored_records
    ]

    max_score = results[0].bm25_score if results else 0.0
    logger.info(
        "ebook_bm25_search_complete corpus=%s candidates=%s max_score=%.6f",
        len(corpus.records),
        len(results),
        max_score,
    )
    return results


def reciprocal_rank_fusion(
    vector_results: list[SearchResult],
    lexical_results: list[SearchResult],
    rank_constant: int,
) -> list[SearchResult]:
    """Fuse vector and lexical rankings with Reciprocal Rank Fusion."""
    by_chunk: dict[int, SearchResult] = {}
    scores: defaultdict[int, float] = defaultdict(float)
    vector_scores: dict[int, float] = {}
    bm25_scores: dict[int, float] = {}

    for rank, result in enumerate(vector_results, start=1):
        by_chunk.setdefault(result.chunk_id, result)
        vector_scores[result.chunk_id] = result.vector_score if result.vector_score is not None else result.score
        scores[result.chunk_id] += 1 / (rank_constant + rank)

    for rank, result in enumerate(lexical_results, start=1):
        by_chunk.setdefault(result.chunk_id, result)
        bm25_scores[result.chunk_id] = result.bm25_score if result.bm25_score is not None else result.score
        scores[result.chunk_id] += 1 / (rank_constant + rank)

    return sorted(
        (
            replace(
                result,
                score=scores[result.chunk_id],
                vector_score=vector_scores.get(result.chunk_id),
                bm25_score=bm25_scores.get(result.chunk_id),
                fused_score=scores[result.chunk_id],
                rank_source="Hybrid",
            )
            for result in by_chunk.values()
        ),
        key=lambda result: result.score,
        reverse=True,
    )


def search_result_from_row(row: Mapping[str, object]) -> SearchResult:
    """Convert a database row mapping into a search result."""
    source_id = row.get("source_id")
    return SearchResult(
        chunk_id=int(row["chunk_id"]),
        text=str(row["text"]),
        source_id=int(source_id) if source_id is not None else None,
        source_title=str(row["source_title"]),
        source_author=optional_str(row["source_author"]),
        chapter_title=optional_str(row["chapter_title"]),
        page_label=optional_str(row["page_label"]),
        score=float(row["score"]) if "score" in row else 0.0,
        vector_score=float(row["score"]) if "score" in row else None,
    )


def optional_str(value: object) -> str | None:
    """Convert nullable database values to optional strings."""
    if value is None:
        return None
    return str(value)


TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")


def tokens(text_value: str) -> list[str]:
    """Extract tokens from a text value.

    This is a simple approximation of the tokenization used by PostgreSQL's full-text search,
    which is sufficient for BM25 candidate retrieval. It lowercases tokens and includes alphanumeric characters and
    underscores.
    """
    return [match.group(0).lower() for match in TOKEN_RE.finditer(text_value)]


QUERY_STOP_WORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "does",
    "for",
    "in",
    "is",
    "of",
    "the",
    "to",
    "what",
    "when",
    "where",
    "which",
    "who",
    "why",
}


def retrieval_query_from_text(query: str) -> str:
    """Remove generic question words while preserving entity and series terms."""
    keywords = [token for token in tokens(query) if token not in QUERY_STOP_WORDS]
    if not keywords:
        return query
    return " ".join(keywords)