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

from __future__ import annotations

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

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

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 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_id,
    new_candidate_row,
    prune_unstorable_unjudged_candidate_phrases,
)
from python.orm.common import get_async_postgres_engine
from python.orm.richie import EbookSource

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncEngine

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

logger = logging.getLogger(__name__)


class BookHasNoChaptersError(ValueError):
    """Raised when a book has no indexed chapter text to generate phrases from."""


async def generate_candidate_phrases_for_books(
    engine: AsyncEngine,
    config: EbookSearchConfig,
) -> PhraseCandidateGenerationResult:
    """Create or refresh candidate phrases for indexed books without calling the LLM judge.

    Every book is submitted to the shared process pool up front and runs in parallel across the
    pool's workers; the call blocks until all books have finished. Each worker opens its own
    database engine from environment variables, loads the book's chapters, and commits the
    book's candidates independently.

    Args:
        engine (AsyncEngine): Engine used to read the book list in this process.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored.

    Results are collected in book order while the pool keeps working. A book failure (including
    a book with no indexed chapters) is logged and counted as not built; the remaining books
    are unaffected.
    """
    async with AsyncSession(engine, expire_on_commit=False) as session:
        source_query = select(EbookSource.id).order_by(EbookSource.id)
        source_ids = (await session.scalars(source_query)).all()
    books_seen = len(source_ids)
    logger.info(
        f"ebook_candidate_phrase_generation_start {books_seen=} {config.phrase_min_tokens=} "
        f"{config.phrase_max_tokens=} {config.protected_phrase_max_candidates_per_book=}"
    )

    pool = get_extraction_pool(config.protected_phrase_extraction_workers)
    wrapped_futures = [
        (
            source_id,
            asyncio.wrap_future(pool.submit(generate_candidate_phrases_for_book_in_worker, source_id, None, config)),
        )
        for source_id in source_ids
    ]
    outcomes: list[BookCandidateResult] = []
    for source_id, wrapped_future in wrapped_futures:
        await asyncio.wait([wrapped_future])
        exception = wrapped_future.exception()
        if exception is not None:
            logger.error(f"ebook_candidate_phrase_generation_book_failed {source_id=}")
            outcomes.append(BookCandidateResult())
            continue
        saved_count = wrapped_future.result()
        logger.info(f"ebook_candidate_phrase_generation_book_committed {source_id=} {saved_count=}")
        outcomes.append(BookCandidateResult(candidates=saved_count, built=True))

    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(
        f"ebook_candidate_phrase_generation_complete {result.books_seen=} {result.books_built=} "
        f"{result.candidate_phrases=}"
    )
    return result


async def recalculate_candidate_phrases_for_book(
    session: AsyncSession,
    source: EbookSource,
    config: EbookSearchConfig,
) -> PhraseRecalculationResult:
    """Remove all book phrase data, regenerate candidates, and commit the completed book.

    Args:
        session (AsyncSession): Active database session; deletion and regeneration commit on it.
        source (EbookSource): Indexed book to recalculate.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated.

    Raises:
        BookHasNoChaptersError: If the book has no indexed chapters. The deletion is rolled
            back, so the book's existing phrases stay intact.

    The deletion and regeneration share the caller's session, so they commit together; a
    regeneration failure rolls the deletion back.
    """
    started_at = perf_counter()
    logger.info(f"ebook_candidate_phrase_recalculation_start {source.id=} {source.title=}")
    deleted = await delete_phrase_data_for_book(session, source.id)
    candidate_count = await generate_candidate_phrases_for_book(
        session,
        source.id,
        series_id=None,
        config=config,
        replace_all=True,
    )

    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(
        f"ebook_candidate_phrase_recalculation_complete {source.id=} {result.deleted_candidates=} "
        f"{result.deleted_protected_phrases=} {result.deleted_aliases=} {result.deleted_mentions=} "
        f"{result.candidate_phrases=} duration_ms={(perf_counter() - started_at) * 1000:.1f}"
    )
    return result


def generate_candidate_phrases_for_book_in_worker(
    book_id: int,
    series_id: int | None,
    config: EbookSearchConfig,
) -> int:
    """Run one book's candidate generation in a pooled worker process.

    The worker has no engine or session to inherit (neither can cross process boundaries), so
    it creates its own engine from environment variables, opens the book's session on it, and
    disposes the engine once the book is stored.

    Args:
        book_id (int): Book the candidates belong to.
        series_id (int | None): Series scope for the stored candidates.
        config (EbookSearchConfig): Runtime phrase-tuning settings.

    Returns:
        int: Number of candidate phrase rows stored.
    """

    async def generate_with_worker_engine() -> int:
        engine = get_async_postgres_engine(name="RICHIE", vector_engine=True, pool_size=1)
        try:
            async with AsyncSession(engine, expire_on_commit=False) as session:
                return await generate_candidate_phrases_for_book(
                    session,
                    book_id,
                    series_id,
                    config,
                )
        finally:
            await engine.dispose()

    return asyncio.run(generate_with_worker_engine())


async def generate_candidate_phrases_for_book(
    session: AsyncSession,
    book_id: int,
    series_id: int | None,
    config: EbookSearchConfig,
    *,
    replace_all: bool = False,
) -> int:
    """Load a book's chapters and metadata, extract candidate phrases, and store them without LLM judging.

    The session commits only when the whole book succeeds; any failure rolls the session back,
    which also restores rows the caller deleted in the same transaction (e.g. a recalculation).

    Args:
        session (AsyncSession): Active database session; committed on success, rolled back on failure.
        book_id (int): Book the candidates belong to.
        series_id (int | None): Series scope for the stored candidates.
        config (EbookSearchConfig): Runtime phrase-tuning settings.
        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.

    Returns:
        int: Number of candidate phrase rows stored.

    Raises:
        BookHasNoChaptersError: If the book has no indexed chapter text.
    """
    started_at = perf_counter()
    chapters = await load_book_chapter_texts(session, book_id)
    if not chapters:
        await session.rollback()
        message = f"book {book_id} has no indexed chapters"
        raise BookHasNoChaptersError(message)
    metadata = await metadata_for_source_id(session, book_id)
    try:
        book_text = "\n\n".join(chapters)
        candidates = extract_phrase_candidates_for_book(
            book_text,
            chapters,
            config,
            metadata=metadata,
        )
        saved_count = await store_candidate_phrases_for_book(
            session,
            book_id,
            series_id,
            candidates,
            config,
            replace_all=replace_all,
        )
        await session.commit()
    except Exception:
        await session.rollback()
        raise
    logger.info(
        f"ebook_candidate_phrase_generation_book_duration {book_id=} {saved_count=} "
        f"duration_ms={(perf_counter() - started_at) * 1000:.1f}"
    )
    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(
            f"ebook_candidate_phrase_save_start {book_id=} candidates={len(limited_candidates)} mode=bulk_insert"
        )
    else:
        pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
        logger.info(
            f"ebook_candidate_phrase_save_start {book_id=} candidates={len(limited_candidates)} {pruned_count=}"
        )
        saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
    logger.info(
        f"ebook_candidate_phrase_save_complete {book_id=} {saved_count=} "
        f"save_ms={(perf_counter() - save_started_at) * 1000:.1f}"
    )
    return saved_count