"""EPUB search models."""

from __future__ import annotations

from datetime import datetime

from pgvector.sqlalchemy import Vector
from sqlalchemy import (
    JSON,
    BigInteger,
    DateTime,
    ForeignKey,
    Index,
    String,
    Text,
    UniqueConstraint,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship

from python.orm.richie.base import TableBase, TableBaseBig

JSON_DOCUMENT = JSON().with_variant(JSONB, "postgresql")


class EbookSource(TableBase):
    """One indexed EPUB file."""

    __tablename__ = "ebook_source"
    __table_args__ = (
        UniqueConstraint("file_path"),
        UniqueConstraint("file_sha256"),
    )

    title: Mapped[str]
    author: Mapped[str | None]
    language: Mapped[str | None]
    publisher: Mapped[str | None]
    identifier: Mapped[str | None]
    file_path: Mapped[str]
    file_sha256: Mapped[str] = mapped_column(String(64))
    file_mtime: Mapped[datetime] = mapped_column(DateTime(timezone=True))
    file_size: Mapped[int] = mapped_column(BigInteger)

    chapters: Mapped[list[EbookChapter]] = relationship(
        "EbookChapter",
        back_populates="source",
        cascade="all, delete-orphan",
        passive_deletes=True,
    )
    chunks: Mapped[list[EbookChunk]] = relationship(
        "EbookChunk",
        back_populates="source",
        cascade="all, delete-orphan",
        passive_deletes=True,
    )


class EbookChapter(TableBase):
    """A chapter or spine document inside an EPUB."""

    __tablename__ = "ebook_chapter"
    __table_args__ = (UniqueConstraint("source_id", "spine_index"),)

    source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
    spine_index: Mapped[int]
    title: Mapped[str | None]
    href: Mapped[str | None]

    source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chapters")
    chunks: Mapped[list[EbookChunk]] = relationship(
        "EbookChunk",
        back_populates="chapter",
        cascade="all, delete-orphan",
        passive_deletes=True,
    )


class EbookChunk(TableBaseBig):
    """A searchable text chunk."""

    __tablename__ = "ebook_chunk"
    __table_args__ = (
        UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"),
        UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"),
    )

    source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
    chapter_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_chapter.id", ondelete="SET NULL"))
    chunk_index: Mapped[int]
    text: Mapped[str]
    token_start: Mapped[int]
    token_count: Mapped[int]
    page_label: Mapped[str | None]
    content_sha256: Mapped[str] = mapped_column(String(64))
    search_text: Mapped[str]

    source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chunks")
    chapter: Mapped[EbookChapter | None] = relationship("EbookChapter", back_populates="chunks")


class EbookEmbeddingModel(TableBase):
    """A supported embedding model."""

    __tablename__ = "ebook_embedding_model"

    name: Mapped[str] = mapped_column(String, unique=True)
    dimension: Mapped[int]
    is_default: Mapped[bool] = mapped_column(default=False)


class EbookChunkEmbedding1024(TableBaseBig):
    """1024-dimensional chunk embedding."""

    __tablename__ = "ebook_chunk_embedding_1024"
    __table_args__ = (
        UniqueConstraint("chunk_id", "model_id"),
        Index(
            "ix_ebook_chunk_embedding_1024_embedding_cosine",
            "embedding",
            postgresql_using="hnsw",
            postgresql_ops={"embedding": "vector_cosine_ops"},
        ),
    )

    chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
    model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
    embedding: Mapped[list[float]] = mapped_column(Vector(1024))


class EbookChunkEmbedding2560(TableBaseBig):
    """2560-dimensional chunk embedding."""

    __tablename__ = "ebook_chunk_embedding_2560"
    __table_args__ = (UniqueConstraint("chunk_id", "model_id"),)

    chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
    model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
    embedding: Mapped[list[float]] = mapped_column(Vector(2560))


class EbookChunkEmbedding4096(TableBaseBig):
    """4096-dimensional chunk embedding."""

    __tablename__ = "ebook_chunk_embedding_4096"
    __table_args__ = (UniqueConstraint("chunk_id", "model_id"),)

    chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
    model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
    embedding: Mapped[list[float]] = mapped_column(Vector(4096))


class EbookCandidatePhrase(TableBase):
    """A high-recall phrase candidate extracted from one book."""

    __tablename__ = "candidate_phrases"
    __table_args__ = (
        UniqueConstraint("book_id", "phrase_norm", name="uq_candidate_phrases_book_id_phrase_norm"),
        Index("candidate_phrases_book_score_idx", "book_id", "candidate_score"),
        Index("candidate_phrases_book_norm_idx", "book_id", "phrase_norm"),
    )

    book_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
    series_id: Mapped[int | None]
    phrase_text: Mapped[str] = mapped_column(Text)
    phrase_norm: Mapped[str] = mapped_column(Text)
    token_count: Mapped[int]
    source_raw_ngram: Mapped[bool] = mapped_column(default=False)
    source_yake: Mapped[bool] = mapped_column(default=False)
    source_spacy_ner: Mapped[bool] = mapped_column(default=False)
    source_spacy_noun_chunk: Mapped[bool] = mapped_column(default=False)
    source_capitalized: Mapped[bool] = mapped_column(default=False)
    source_metadata: Mapped[bool] = mapped_column(default=False)
    spacy_label: Mapped[str | None]
    raw_count: Mapped[int] = mapped_column(default=0)
    chapter_count: Mapped[int] = mapped_column(default=0)
    yake_score: Mapped[float | None]
    candidate_score: Mapped[float] = mapped_column(default=0.0)
    sample_contexts: Mapped[list[str] | None] = mapped_column(JSON_DOCUMENT)
    llm_judged: Mapped[bool] = mapped_column(default=False)
    llm_keep: Mapped[bool | None]
    llm_confidence: Mapped[float | None]
    llm_category: Mapped[str | None]
    llm_reason: Mapped[str | None] = mapped_column(Text)


class EbookProtectedPhrase(TableBase):
    """A phrase accepted by the LLM judge for protected query matching."""

    __tablename__ = "protected_phrases"
    __table_args__ = (
        UniqueConstraint("book_id", "phrase_norm", name="uq_protected_phrases_book_id_phrase_norm"),
        Index("protected_phrases_norm_idx", "phrase_norm"),
        Index("protected_phrases_book_norm_idx", "book_id", "phrase_norm"),
        Index("protected_phrases_series_norm_idx", "series_id", "phrase_norm"),
    )

    book_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
    series_id: Mapped[int | None]
    phrase_text: Mapped[str] = mapped_column(Text)
    phrase_norm: Mapped[str] = mapped_column(Text)
    canonical_id: Mapped[str]
    phrase_type: Mapped[str | None]
    token_count: Mapped[int]
    confidence: Mapped[float]
    importance: Mapped[float] = mapped_column(default=0.5)
    allow_nested: Mapped[bool] = mapped_column(default=False)
    suppress_children: Mapped[bool] = mapped_column(default=True)
    source_candidate_id: Mapped[int | None] = mapped_column(
        ForeignKey("main.candidate_phrases.id", ondelete="SET NULL")
    )


class EbookPhraseAlias(TableBase):
    """A normalized alias that maps to a protected phrase."""

    __tablename__ = "phrase_aliases"
    __table_args__ = (
        UniqueConstraint("phrase_id", "alias_norm", name="uq_phrase_aliases_phrase_id_alias_norm"),
        Index("phrase_aliases_norm_idx", "alias_norm"),
    )

    phrase_id: Mapped[int] = mapped_column(ForeignKey("main.protected_phrases.id", ondelete="CASCADE"))
    alias_text: Mapped[str] = mapped_column(Text)
    alias_norm: Mapped[str] = mapped_column(Text)
    confidence: Mapped[float] = mapped_column(default=1.0)


class EbookChunkPhraseMention(TableBase):
    """A precomputed occurrence of a protected phrase inside one chunk."""

    __tablename__ = "chunk_phrase_mentions"
    __table_args__ = (
        UniqueConstraint("chunk_id", "phrase_id", "start_char", name="uq_chunk_phrase_mentions_chunk_phrase_start"),
        Index("chunk_phrase_mentions_phrase_idx", "phrase_id"),
        Index("chunk_phrase_mentions_chunk_idx", "chunk_id"),
    )

    chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
    phrase_id: Mapped[int] = mapped_column(ForeignKey("main.protected_phrases.id", ondelete="CASCADE"))
    book_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
    series_id: Mapped[int | None]
    start_char: Mapped[int]
    end_char: Mapped[int | None]