"""Database persistence for candidate and protected phrase rows."""

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING

from sqlalchemy import delete, func, or_, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert

from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
from python.ebook_search.protected_phrases.models import (
    CorpusPhraseStats,
    PhraseCandidate,
    PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.text_normalization import normalize_text
from python.orm.richie import (
    EbookCandidatePhrase,
    EbookChunk,
    EbookChunkPhraseMention,
    EbookPhraseAlias,
    EbookProtectedPhrase,
    EbookSource,
)

if TYPE_CHECKING:
    from collections.abc import Sequence

    from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
    from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
    from sqlalchemy.ext.asyncio import AsyncSession

    from python.ebook_search.config import EbookSearchConfig
    from python.ebook_search.protected_phrases.models import LLMJudgment
    from python.orm.richie.base import TableBase

logger = logging.getLogger(__name__)


def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
    """Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.

    Production runs on PostgreSQL while tests run on SQLite; both support upserts with
    compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.

    Args:
        session (AsyncSession): Active database session whose bind selects the dialect.
        table (type[TableBase]): Mapped table to insert into.

    Returns:
        PostgresInsert | SqliteInsert: A dialect insert exposing ``on_conflict_do_update``.
    """
    if session.get_bind().dialect.name == "sqlite":
        return sqlite_insert(table)
    return pg_insert(table)


async def load_book_text(session: AsyncSession, book_id: int) -> str:
    """Load a book's indexed chunk text as one string for phrase extraction.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose chunk text is loaded.

    Returns:
        str: The book's chunk text joined into a single string.
    """
    texts = await session.scalars(
        select(EbookChunk.text).where(EbookChunk.source_id == book_id).order_by(EbookChunk.chunk_index)
    )
    return "\n\n".join(stripped for text in texts if (stripped := text.strip()))


async def load_book_chapter_texts(session: AsyncSession, book_id: int) -> list[str]:
    """Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose chunks are grouped into chapters.

    Returns:
        list[str]: Non-empty chapter-like text blocks in chunk order.
    """
    rows = await session.execute(
        select(EbookChunk.chapter_id, EbookChunk.text)
        .where(EbookChunk.source_id == book_id)
        .order_by(EbookChunk.chunk_index)
    )
    chapters: list[str] = []
    current_chapter_id: int | None = None
    current_parts: list[str] = []
    have_current = False

    for chapter_id, text in rows:
        if have_current and chapter_id != current_chapter_id:
            chapter_text = "\n\n".join(current_parts).strip()
            if chapter_text:
                chapters.append(chapter_text)
            current_parts = []
        current_chapter_id = chapter_id
        current_parts.append(str(text))
        have_current = True

    if current_parts:
        chapter_text = "\n\n".join(current_parts).strip()
        if chapter_text:
            chapters.append(chapter_text)
    return chapters


def metadata_for_source(source: EbookSource) -> dict[str, object | None]:
    """Return phrase extraction metadata for one indexed source.

    Args:
        source (EbookSource): Indexed source to read metadata from.

    Returns:
        dict[str, object | None]: Title, author, language, publisher, and identifier values.
    """
    return {
        "title": source.title,
        "author": source.author,
        "language": source.language,
        "publisher": source.publisher,
        "identifier": source.identifier,
    }


async def metadata_for_source_id(session: AsyncSession, source_id: int) -> dict[str, object | None]:
    """Return phrase extraction metadata for one indexed source by id.

    Args:
        session (AsyncSession): Active database session.
        source_id (int): Id of the indexed source to read metadata from.

    Returns:
        dict[str, object | None]: Title, author, language, publisher, and identifier values.

    Raises:
        ValueError: If no source exists with the given id.
    """
    source = await session.get(EbookSource, source_id)
    if source is None:
        msg = f"No indexed source with id {source_id}"
        raise ValueError(msg)
    return metadata_for_source(source)


async def count_protected_phrases(session: AsyncSession, book_id: int) -> int:
    """Count stored protected phrases for one book.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose protected phrases are counted.

    Returns:
        int: Number of protected phrases stored for the book.
    """
    return (
        await session.scalars(
            select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
        )
    ).one()


async def count_unjudged_candidates(session: AsyncSession, book_id: int, config: EbookSearchConfig) -> int:
    """Count storable candidate rows for a book that have not yet been judged.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose unjudged candidates are counted.
        config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.

    Returns:
        int: Number of storable, unjudged candidate rows for the book.
    """
    return (
        await session.scalars(
            select(func.count(EbookCandidatePhrase.id)).where(
                EbookCandidatePhrase.book_id == book_id,
                EbookCandidatePhrase.llm_judged.is_(False),
                EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
                EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
            )
        )
    ).one()


async def corpus_phrase_stats(session: AsyncSession) -> CorpusPhraseStats:
    """Summarize candidate and protected phrase coverage across the whole corpus.

    Args:
        session (AsyncSession): Active database session.

    Returns:
        CorpusPhraseStats: Corpus-wide phrase counts and per-book coverage counts.
    """
    total_books = (await session.scalars(select(func.count(EbookSource.id)))).one()
    candidate_phrases, judged_candidates, books_with_candidates, books_with_unjudged = (
        await session.execute(
            select(
                func.count(EbookCandidatePhrase.id),
                func.count(EbookCandidatePhrase.id).filter(EbookCandidatePhrase.llm_judged.is_(True)),
                func.count(func.distinct(EbookCandidatePhrase.book_id)),
                func.count(func.distinct(EbookCandidatePhrase.book_id)).filter(
                    EbookCandidatePhrase.llm_judged.is_(False)
                ),
            )
        )
    ).one()
    protected_phrases = (await session.scalars(select(func.count(EbookProtectedPhrase.id)))).one()
    return CorpusPhraseStats(
        total_books=total_books,
        books_with_candidates=books_with_candidates,
        books_fully_judged=books_with_candidates - books_with_unjudged,
        candidate_phrases=candidate_phrases,
        judged_candidates=judged_candidates,
        unjudged_candidates=candidate_phrases - judged_candidates,
        protected_phrases=protected_phrases,
    )


async def book_ids_pending_first_judgment(session: AsyncSession) -> list[int]:
    """Return books that have candidate phrases but no judged candidates yet.

    Args:
        session (AsyncSession): Active database session.

    Returns:
        list[int]: Book ids with candidates where judging has never run, ordered by id.
    """
    judged_books = select(EbookCandidatePhrase.book_id).where(EbookCandidatePhrase.llm_judged.is_(True)).distinct()
    return list(
        (
            await session.scalars(
                select(EbookCandidatePhrase.book_id)
                .where(EbookCandidatePhrase.book_id.not_in(judged_books))
                .distinct()
                .order_by(EbookCandidatePhrase.book_id)
            )
        ).all()
    )


async def load_candidates_for_judgment(
    session: AsyncSession,
    book_id: int,
    config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
    """Load every storable unjudged candidate row for a book.

    Rows may have been stored before the current junk filters and score weights existed, so
    callers re-check :func:`is_junk_phrase` and rescore before selecting what to judge.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose candidates are loaded.
        config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.

    Returns:
        Sequence[EbookCandidatePhrase]: Storable, unjudged candidate rows ordered by stored score.
    """
    query = (
        select(EbookCandidatePhrase)
        .where(
            EbookCandidatePhrase.book_id == book_id,
            EbookCandidatePhrase.llm_judged.is_(False),
            EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
            EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
        )
        .order_by(
            EbookCandidatePhrase.candidate_score.desc(),
            EbookCandidatePhrase.raw_count.desc(),
            EbookCandidatePhrase.id,
        )
    )
    return (await session.scalars(query)).all()


def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
    """Recreate an in-memory candidate from a persisted candidate row.

    Args:
        row (EbookCandidatePhrase): Stored candidate row to convert.

    Returns:
        PhraseCandidate: An in-memory candidate mirroring the row's fields.
    """
    return PhraseCandidate(
        phrase_text=row.phrase_text,
        phrase_norm=row.phrase_norm,
        token_count=row.token_count,
        source_raw_ngram=row.source_raw_ngram,
        source_yake=row.source_yake,
        source_spacy_ner=row.source_spacy_ner,
        source_spacy_noun_chunk=row.source_spacy_noun_chunk,
        source_capitalized=row.source_capitalized,
        source_metadata=row.source_metadata,
        spacy_label=row.spacy_label,
        raw_count=row.raw_count,
        chapter_count=row.chapter_count,
        yake_score=row.yake_score,
        candidate_score=row.candidate_score,
        sample_contexts=row.sample_contexts or [],
    )


def candidate_row_values(
    book_id: int,
    series_id: int | None,
    candidate: PhraseCandidate,
    *,
    judgment: LLMJudgment | None,
) -> dict[str, object]:
    """Build the column values for one candidate phrase upsert.

    Args:
        book_id (int): Book the candidate belongs to.
        series_id (int | None): Series scope stored on the row.
        candidate (PhraseCandidate): Candidate whose fields are written to the row.
        judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.

    Returns:
        dict[str, object]: Column values keyed by column name.
    """
    values: dict[str, object] = {
        "book_id": book_id,
        "phrase_norm": candidate.phrase_norm,
        "series_id": series_id,
        "phrase_text": candidate.phrase_text,
        "token_count": candidate.token_count,
        "source_raw_ngram": candidate.source_raw_ngram,
        "source_yake": candidate.source_yake,
        "source_spacy_ner": candidate.source_spacy_ner,
        "source_spacy_noun_chunk": candidate.source_spacy_noun_chunk,
        "source_capitalized": candidate.source_capitalized,
        "source_metadata": candidate.source_metadata,
        "spacy_label": candidate.spacy_label,
        "raw_count": candidate.raw_count,
        "chapter_count": candidate.chapter_count,
        "yake_score": candidate.yake_score,
        "candidate_score": candidate.candidate_score,
        "llm_judged": judgment is not None,
    }
    if candidate.sample_contexts:
        values["sample_contexts"] = list(candidate.sample_contexts)
    if judgment is not None:
        values.update(
            llm_keep=judgment.keep,
            llm_confidence=judgment.confidence,
            llm_category=judgment.category,
            llm_reason=judgment.reason,
        )
    return values


async def save_candidate_to_db(
    session: AsyncSession,
    book_id: int,
    series_id: int | None,
    candidate: PhraseCandidate,
    *,
    judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
    """Insert or update one candidate phrase row.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book the candidate belongs to.
        series_id (int | None): Series scope stored on the row.
        candidate (PhraseCandidate): Candidate whose fields are written to the row.
        judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.

    Returns:
        EbookCandidatePhrase: The inserted or updated candidate row.
    """
    values = candidate_row_values(book_id, series_id, candidate, judgment=judgment)

    # Preserve an existing judgment when this call is only refreshing candidate fields.
    skip_update = {"book_id", "phrase_norm"}
    if judgment is None:
        skip_update.add("llm_judged")
    insert_statement = dialect_insert(session, EbookCandidatePhrase).values(**values)
    statement = insert_statement.on_conflict_do_update(
        index_elements=["book_id", "phrase_norm"],
        set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
    ).returning(EbookCandidatePhrase)
    return (await session.scalars(statement, execution_options={"populate_existing": True})).one()


BULK_CANDIDATE_UPSERT_CHUNK = 1000


async def bulk_upsert_unjudged_candidates(
    session: AsyncSession,
    book_id: int,
    series_id: int | None,
    candidates: Sequence[PhraseCandidate],
) -> int:
    """Insert or update many freshly extracted candidate rows in chunked multi-row upserts.

    Saving one row per statement costs one database round trip per candidate, which dominated
    generation time for full books, so candidates are written ``BULK_CANDIDATE_UPSERT_CHUNK``
    rows per statement instead. Existing judgments and sample contexts are never overwritten:
    fresh extractions carry no contexts, and ``llm_judged`` plus the ``llm_*`` columns are left
    out of the conflict update. Candidates must have unique ``phrase_norm`` values, as produced
    by extraction, since one multi-row upsert cannot touch the same row twice.

    Args:
        session (Session): Active database session.
        book_id (int): Book the candidates belong to.
        series_id (int | None): Series scope stored on the rows.
        candidates (Sequence[PhraseCandidate]): Freshly extracted candidates to persist.

    Returns:
        int: Number of candidate rows written.
    """
    values = [
        candidate_row_values(book_id, series_id, candidate, judgment=None)
        for candidate in candidates
        if not candidate.sample_contexts
    ]
    if len(values) != len(candidates):
        msg = "bulk_upsert_unjudged_candidates only accepts freshly extracted candidates without sample contexts"
        raise ValueError(msg)
    skip_update = {"book_id", "phrase_norm", "llm_judged"}
    for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
        chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
        insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
        statement = insert_statement.on_conflict_do_update(
            index_elements=["book_id", "phrase_norm"],
            set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
        )
        await session.execute(statement)
    return len(values)


def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCandidate) -> EbookCandidatePhrase:
    """Build a fresh unjudged candidate row without checking for an existing one.

    Unlike :func:`save_candidate_to_db`, this does no lookup, so it is only safe when the caller
    guarantees there is no existing row for ``(book_id, candidate.phrase_norm)`` — for example
    right after :func:`delete_phrase_data_for_book` has cleared the book.

    Args:
        book_id (int): Book the candidate belongs to.
        series_id (int | None): Series scope stored on the row.
        candidate (PhraseCandidate): Candidate whose fields are written to the row.

    Returns:
        EbookCandidatePhrase: A new, unattached candidate row.
    """
    row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm)
    row.llm_judged = False
    row.series_id = series_id
    row.phrase_text = candidate.phrase_text
    row.token_count = candidate.token_count
    row.source_raw_ngram = candidate.source_raw_ngram
    row.source_yake = candidate.source_yake
    row.source_spacy_ner = candidate.source_spacy_ner
    row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk
    row.source_capitalized = candidate.source_capitalized
    row.source_metadata = candidate.source_metadata
    row.spacy_label = candidate.spacy_label
    row.raw_count = candidate.raw_count
    row.chapter_count = candidate.chapter_count
    row.yake_score = candidate.yake_score
    row.candidate_score = candidate.candidate_score
    if candidate.sample_contexts:
        row.sample_contexts = list(candidate.sample_contexts)
    return row


async def upsert_protected_phrase(
    session: AsyncSession,
    book_id: int,
    series_id: int | None,
    candidate: PhraseCandidate,
    judgment: LLMJudgment,
    source_candidate: EbookCandidatePhrase,
) -> EbookProtectedPhrase:
    """Insert or update one accepted protected phrase and its aliases.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book the protected phrase belongs to.
        series_id (int | None): Series scope stored on the phrase.
        candidate (PhraseCandidate): Candidate the phrase was promoted from.
        judgment (LLMJudgment): Accepted judgment supplying canonical text, category, and aliases.
        source_candidate (EbookCandidatePhrase): Candidate row the phrase was promoted from.

    Returns:
        EbookProtectedPhrase: The inserted or updated protected phrase row.

    Raises:
        ValueError: If the chosen phrase text normalizes to empty.
    """
    phrase_text = judgment.canonical or candidate.phrase_text
    phrase_norm = normalize_text(phrase_text)
    if not phrase_norm:
        msg = f"Protected phrase normalized to empty text: {phrase_text!r}"
        raise ValueError(msg)

    values = {
        "book_id": book_id,
        "phrase_norm": phrase_norm,
        "series_id": series_id,
        "phrase_text": phrase_text,
        "canonical_id": make_canonical_id(judgment, phrase_norm),
        "phrase_type": judgment.category,
        "token_count": len(phrase_norm.split()),
        "confidence": judgment.confidence,
        "importance": judgment.importance,
        "allow_nested": judgment.allow_nested,
        "suppress_children": judgment.suppress_children,
        "source_candidate_id": source_candidate.id,
    }
    insert_statement = dialect_insert(session, EbookProtectedPhrase).values(**values)
    statement = insert_statement.on_conflict_do_update(
        index_elements=["book_id", "phrase_norm"],
        set_={
            column: insert_statement.excluded[column] for column in values if column not in {"book_id", "phrase_norm"}
        },
    ).returning(EbookProtectedPhrase)
    row = (await session.scalars(statement, execution_options={"populate_existing": True})).one()

    for alias_text in judgment.aliases:
        await upsert_phrase_alias(session, row, alias_text)
    return row


async def upsert_phrase_alias(
    session: AsyncSession,
    phrase: EbookProtectedPhrase,
    alias_text: str,
) -> EbookPhraseAlias | None:
    """Insert or update one protected phrase alias.

    Args:
        session (AsyncSession): Active database session.
        phrase (EbookProtectedPhrase): Protected phrase the alias points to.
        alias_text (str): Alias surface form to store.

    Returns:
        EbookPhraseAlias | None: The alias row, or ``None`` when the alias is empty or equals the phrase.
    """
    alias_norm = normalize_text(alias_text)
    if not alias_norm or alias_norm == phrase.phrase_norm:
        return None

    insert_statement = dialect_insert(session, EbookPhraseAlias).values(
        phrase_id=phrase.id,
        alias_norm=alias_norm,
        alias_text=alias_text,
        confidence=1.0,
    )
    statement = insert_statement.on_conflict_do_update(
        index_elements=["phrase_id", "alias_norm"],
        set_={
            "alias_text": insert_statement.excluded.alias_text,
            "confidence": insert_statement.excluded.confidence,
        },
    ).returning(EbookPhraseAlias)
    return (await session.scalars(statement, execution_options={"populate_existing": True})).one()


def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str:
    """Create a deterministic canonical id from a judgment category and phrase.

    Args:
        judgment (LLMJudgment): Judgment supplying the phrase category.
        phrase_norm (str): Normalized phrase text to slugify.

    Returns:
        str: A ``category:slug`` canonical identifier.
    """
    category = slugify_identifier(judgment.category or "phrase")
    phrase_slug = slugify_identifier(phrase_norm)
    return f"{category}:{phrase_slug}"


def slugify_identifier(value: str) -> str:
    """Normalize text for use inside a canonical id.

    Args:
        value (str): Text to slugify.

    Returns:
        str: A lowercase underscore slug, or ``"unknown"`` when empty.
    """
    slug = re.sub(r"[^a-z0-9]+", "_", normalize_text(value).replace("'", ""))
    return slug.strip("_") or "unknown"


async def prune_unstorable_unjudged_candidate_phrases(
    session: AsyncSession,
    book_id: int,
    config: EbookSearchConfig,
) -> int:
    """Delete old unjudged candidate rows that no longer satisfy storage filters.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose stale candidates are pruned.
        config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.

    Returns:
        int: Number of candidate rows deleted.
    """
    deleted = rowcount(
        await session.execute(
            delete(EbookCandidatePhrase).where(
                EbookCandidatePhrase.book_id == book_id,
                EbookCandidatePhrase.llm_judged.is_(False),
                or_(
                    EbookCandidatePhrase.token_count < config.phrase_min_tokens,
                    EbookCandidatePhrase.raw_count < minimum_candidate_raw_count(config),
                ),
            )
        )
    )
    if deleted:
        logger.info(
            "ebook_candidate_phrase_unstorable_pruned book_id=%s deleted=%s min_tokens=%s min_uses=%s",
            book_id,
            deleted,
            config.phrase_min_tokens,
            minimum_candidate_raw_count(config),
        )
    return deleted


async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> PhraseRecalculationResult:
    """Delete all candidate, protected, alias, and mention phrase data for one book.

    Args:
        session (AsyncSession): Active database session.
        book_id (int): Book whose phrase data is deleted.

    Returns:
        PhraseRecalculationResult: Deleted-row counts with ``candidate_phrases`` set to 0.
    """
    protected_ids = (
        await session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id))
    ).all()
    deleted_aliases = 0
    if protected_ids:
        deleted_aliases = rowcount(
            await session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
        )

    deleted_mentions = rowcount(
        await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
    )
    if protected_ids:
        deleted_mentions += rowcount(
            await session.execute(
                delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids))
            )
        )

    deleted_protected = rowcount(
        await session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
    )
    deleted_candidates = rowcount(
        await session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
    )
    await session.flush()
    logger.info(
        "ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
        book_id,
        deleted_candidates,
        deleted_protected,
        deleted_aliases,
        deleted_mentions,
    )
    return PhraseRecalculationResult(
        book_id=book_id,
        deleted_candidates=deleted_candidates,
        deleted_protected_phrases=deleted_protected,
        deleted_aliases=deleted_aliases,
        deleted_mentions=deleted_mentions,
        candidate_phrases=0,
    )


def rowcount(result: object) -> int:
    """Return a safe integer rowcount from a SQLAlchemy execution result.

    Args:
        result (object): SQLAlchemy execution result that may expose ``rowcount``.

    Returns:
        int: The result's rowcount, or 0 when it is missing or negative.
    """
    count = getattr(result, "rowcount", 0)
    return int(count if count is not None and count >= 0 else 0)