"""Book-level orchestration for candidate n-gram generation and recalculation."""

from __future__ import annotations

import asyncio
import logging
from collections import deque
from time import perf_counter
from typing import TYPE_CHECKING

from sqlalchemy import select

from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
from python.ebook_search.protected_phrases.models import (
    BookCandidateResult,
    PhraseCandidateGenerationResult,
    PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.pool import extract_phrase_candidates_in_pool, get_extraction_pool
from python.ebook_search.protected_phrases.store import (
    bulk_upsert_unjudged_candidates,
    delete_phrase_data_for_book,
    load_book_chapter_texts,
    metadata_for_source,
    new_candidate_row,
    prune_unstorable_unjudged_candidate_phrases,
)
from python.orm.richie import EbookCandidatePhrase, EbookSource

if TYPE_CHECKING:
    from collections.abc import Mapping, Sequence
    from concurrent.futures import Future

    from sqlalchemy.ext.asyncio import AsyncSession

    from python.ebook_search.config import EbookSearchConfig
    from python.ebook_search.protected_phrases.extraction import SpacyLanguage
    from python.ebook_search.protected_phrases.models import PhraseCandidate

logger = logging.getLogger(__name__)


async def generate_candidate_phrases_for_books(
    session: AsyncSession,
    config: EbookSearchConfig,
    *,
    only_missing: bool = False,
) -> PhraseCandidateGenerationResult:
    """Create or refresh candidate phrases for indexed books without calling the LLM judge.

    Extraction always runs concurrently in the shared process pool so a full backfill uses
    multiple cores.

    Args:
        session (Session): Active database session.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        only_missing (bool): When True, only generate for books that have no candidate phrases
            yet instead of refreshing every book.

    Returns:
        PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored.
    """
    source_query = select(EbookSource).order_by(EbookSource.id)
    if only_missing:
        has_candidates = select(EbookCandidatePhrase.id).where(EbookCandidatePhrase.book_id == EbookSource.id)
        source_query = source_query.where(~has_candidates.exists())
    sources = (await session.scalars(source_query)).all()
    books_seen = len(sources)
    logger.info(
        "ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s",
        books_seen,
        config.phrase_min_tokens,
        config.phrase_max_tokens,
        config.protected_phrase_max_candidates_per_book,
    )

    outcomes = await generate_candidates_for_sources_pooled(session, sources, config)

    result = PhraseCandidateGenerationResult(
        books_seen=books_seen,
        books_built=sum(1 for outcome in outcomes if outcome.built),
        candidate_phrases=sum(outcome.candidates for outcome in outcomes),
    )
    logger.info(
        "ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s",
        result.books_seen,
        result.books_built,
        result.candidate_phrases,
    )
    return result


async def generate_candidates_for_sources_pooled(
    session: AsyncSession,
    sources: Sequence[EbookSource],
    config: EbookSearchConfig,
) -> list[BookCandidateResult]:
    """Generate candidate phrases for many books, extracting them concurrently in worker processes.

    Chapter loading and row persistence stay on the caller's session (serial), while the CPU-bound
    extraction runs in the shared process pool. A bounded window of in-flight books overlaps
    extraction across cores without loading every book's candidates into memory at once.

    Args:
        session (Session): Active database session.
        sources (Sequence[EbookSource]): Indexed books to generate candidates for.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        list[BookCandidateResult]: One result per book.
    """
    pool = get_extraction_pool(config.protected_phrase_extraction_workers)
    max_in_flight = max(1, config.protected_phrase_extraction_workers) * 2
    pending: deque[tuple[EbookSource, Future[list[PhraseCandidate]]]] = deque()
    outcomes: list[BookCandidateResult] = []

    async def drain_one() -> None:
        source, future = pending.popleft()
        extracted = await asyncio.wrap_future(future)
        outcomes.append(await store_source_candidates(session, source, extracted, config))

    try:
        for source in sources:
            chapters = await load_book_chapter_texts(session, source.id)
            if not chapters:
                logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id)
                outcomes.append(BookCandidateResult())
                continue
            future = pool.submit(
                extract_phrase_candidates_for_book,
                "\n\n".join(chapters),
                chapters,
                config,
                metadata=metadata_for_source(source),
            )
            pending.append((source, future))
            if len(pending) >= max_in_flight:
                await drain_one()
        while pending:
            await drain_one()
    except Exception:
        for _, future in pending:
            future.cancel()
        await session.rollback()
        logger.exception("ebook_candidate_phrase_generation_pooled_failed")
        raise
    return outcomes


async def store_source_candidates(
    session: AsyncSession,
    source: EbookSource,
    limited_candidates: list[PhraseCandidate],
    config: EbookSearchConfig,
) -> BookCandidateResult:
    """Persist and commit one book's already-extracted candidates.

    Args:
        session (AsyncSession): Active database session.
        source (EbookSource): Book the candidates belong to.
        limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        BookCandidateResult: Candidate count and that the book was committed.
    """
    book_started_at = perf_counter()
    saved_count = await store_candidate_phrases_for_book(session, source.id, None, limited_candidates, config)
    await session.commit()
    logger.info(
        "ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f",
        source.id,
        saved_count,
        (perf_counter() - book_started_at) * 1000,
    )
    return BookCandidateResult(candidates=saved_count, built=True)


async def recalculate_candidate_phrases_for_book(
    session: AsyncSession,
    source: EbookSource,
    config: EbookSearchConfig,
    *,
    nlp: SpacyLanguage | None = None,
    use_process_pool: bool = False,
) -> PhraseRecalculationResult:
    """Remove all book phrase data, regenerate candidates, and commit the completed book.

    Args:
        session (Session): Active database session.
        source (EbookSource): Indexed book to recalculate.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
        use_process_pool (bool): Run the CPU-bound extraction in a worker process so concurrent
            recalculations do not serialize behind the GIL. Defaults to in-process for callers
            (tests, backfills) that do not need it.

    Returns:
        PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated.
    """
    started_at = perf_counter()
    logger.info(
        "ebook_candidate_phrase_recalculation_start source_id=%s title=%r",
        source.id,
        source.title,
    )
    try:
        deleted = await delete_phrase_data_for_book(session, source.id)
        chapters = await load_book_chapter_texts(session, source.id)
        if not chapters:
            logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id)
            await session.commit()
            return PhraseRecalculationResult(
                book_id=source.id,
                deleted_candidates=deleted.deleted_candidates,
                deleted_protected_phrases=deleted.deleted_protected_phrases,
                deleted_aliases=deleted.deleted_aliases,
                deleted_mentions=deleted.deleted_mentions,
                candidate_phrases=0,
            )

        candidate_count = await generate_candidate_phrases_for_book(
            session,
            source.id,
            series_id=None,
            chapters=chapters,
            config=config,
            nlp=nlp,
            metadata=metadata_for_source(source),
            replace_all=True,
            use_process_pool=use_process_pool,
        )
        await session.commit()
    except Exception:
        await session.rollback()
        logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id)
        raise

    result = PhraseRecalculationResult(
        book_id=source.id,
        deleted_candidates=deleted.deleted_candidates,
        deleted_protected_phrases=deleted.deleted_protected_phrases,
        deleted_aliases=deleted.deleted_aliases,
        deleted_mentions=deleted.deleted_mentions,
        candidate_phrases=candidate_count,
    )
    logger.info(
        "ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s "
        "deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f",
        source.id,
        result.deleted_candidates,
        result.deleted_protected_phrases,
        result.deleted_aliases,
        result.deleted_mentions,
        result.candidate_phrases,
        (perf_counter() - started_at) * 1000,
    )
    return result


async def generate_candidate_phrases_for_book(
    session: AsyncSession,
    book_id: int,
    series_id: int | None,
    chapters: Sequence[str],
    config: EbookSearchConfig,
    *,
    nlp: SpacyLanguage | None = None,
    metadata: Mapping[str, object] | None = None,
    replace_all: bool = False,
    use_process_pool: bool = False,
) -> int:
    """Extract and store candidate phrases for one book without LLM judging.

    Args:
        session (Session): Active database session.
        book_id (int): Book the candidates belong to.
        series_id (int | None): Series scope for the stored candidates.
        chapters (Sequence[str]): Chapter-like text blocks used for extraction and frequency counts.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
        metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
        replace_all (bool): When the caller has already cleared this book's candidates (e.g. a
            recalculation), skip the per-candidate existence lookup and bulk-insert new rows.
        use_process_pool (bool): Run the CPU-bound extraction in a worker process to avoid
            serializing concurrent requests behind the GIL. Ignored when ``nlp`` is set, since
            the spaCy pipeline cannot be sent to a worker process.

    Returns:
        int: Number of candidate phrase rows stored.
    """
    started_at = perf_counter()
    book_text = "\n\n".join(chapters)
    if use_process_pool and nlp is None:
        limited_candidates = await extract_phrase_candidates_in_pool(book_text, chapters, config, metadata=metadata)
    else:
        limited_candidates = extract_phrase_candidates_for_book(
            book_text,
            chapters,
            config,
            nlp=nlp,
            metadata=metadata,
        )
    saved_count = await store_candidate_phrases_for_book(
        session,
        book_id,
        series_id,
        limited_candidates,
        config,
        replace_all=replace_all,
    )
    logger.info(
        "ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f",
        book_id,
        saved_count,
        (perf_counter() - started_at) * 1000,
    )
    return saved_count


async def store_candidate_phrases_for_book(
    session: AsyncSession,
    book_id: int,
    series_id: int | None,
    limited_candidates: list[PhraseCandidate],
    config: EbookSearchConfig,
    *,
    replace_all: bool = False,
) -> int:
    """Persist already-extracted candidate phrase rows for one book without committing.

    Args:
        session (Session): Active database session.
        book_id (int): Book the candidates belong to.
        series_id (int | None): Series scope for the stored candidates.
        limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        replace_all (bool): When the caller has already cleared this book's candidates, skip the
            per-candidate existence lookup and bulk-insert new rows.

    Returns:
        int: Number of candidate phrase rows stored.
    """
    save_started_at = perf_counter()
    if replace_all:
        rows = [new_candidate_row(book_id, series_id, candidate) for candidate in limited_candidates]
        session.add_all(rows)
        await session.flush()
        saved_count = len(rows)
        logger.info(
            "ebook_candidate_phrase_save_start book_id=%s candidates=%s mode=bulk_insert",
            book_id,
            len(limited_candidates),
        )
    else:
        pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
        logger.info(
            "ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s",
            book_id,
            len(limited_candidates),
            pruned_count,
        )
        saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
    logger.info(
        "ebook_candidate_phrase_save_complete book_id=%s candidates=%s save_ms=%.1f",
        book_id,
        saved_count,
        (perf_counter() - save_started_at) * 1000,
    )
    return saved_count