"""SQLAlchemy models for Quber RAG database."""

from datetime import datetime
from typing import Any, Optional

from pgvector.sqlalchemy import Vector
from sqlalchemy import JSON, ForeignKey, Integer, String, Text, UniqueConstraint
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class Base(DeclarativeBase):
    """Base class for all database models."""

    pass


class Document(Base):
    """Represents a processed document with extracted tables."""

    __tablename__ = "documents"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    filename: Mapped[str] = mapped_column(String(255), nullable=False, unique=True, index=True)
    extraction_date: Mapped[datetime] = mapped_column(nullable=False)
    total_pages: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    total_tables: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    executive_summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    model_provider: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
    model: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    created_at: Mapped[datetime] = mapped_column(default=datetime.now)

    # Relationship to extracted tables
    tables: Mapped[list["ExtractedTable"]] = relationship(
        "ExtractedTable", back_populates="document", cascade="all, delete-orphan"
    )

    def __repr__(self) -> str:
        return f"<Document(id={self.id}, filename='{self.filename}', tables={self.total_tables})>"


class ExtractedTable(Base):
    """Represents an extracted table with embeddings for semantic search."""

    __tablename__ = "extracted_tables"
    __table_args__ = (UniqueConstraint("document_id", "table_id", name="uq_document_table"),)

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    document_id: Mapped[int] = mapped_column(ForeignKey("documents.id", ondelete="CASCADE"), index=True)
    table_id: Mapped[int] = mapped_column(Integer, nullable=False)
    page_number: Mapped[int] = mapped_column(Integer, nullable=False)
    # `extractor` is meant to name the engine that produced the row, but no
    # writer sets it. Rows imported by `quber db import-json` and
    # `quber db import-dir` leave it NULL.
    extractor: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
    procedural_title: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    llm_title: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    llm_subtitle: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    llm_description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    table_markdown: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    headers: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True)
    footnotes: Mapped[Optional[list[str]]] = mapped_column(JSON, nullable=True)
    bbox: Mapped[Optional[list[float]]] = mapped_column(JSON, nullable=True)
    flavor: Mapped[Optional[str]] = mapped_column(String(20), nullable=True)
    table_metadata: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True)

    # Vector embeddings (1024 dims for bge-large-en-v1.5)
    title_embedding: Mapped[Optional[Any]] = mapped_column(Vector(1024), nullable=True)
    description_embedding: Mapped[Optional[Any]] = mapped_column(Vector(1024), nullable=True)

    created_at: Mapped[datetime] = mapped_column(default=datetime.now)

    # Relationship to document
    document: Mapped["Document"] = relationship("Document", back_populates="tables")

    def __repr__(self) -> str:
        return (
            f"<ExtractedTable(id={self.id}, document_id={self.document_id}, "
            f"table_id={self.table_id}, page={self.page_number})>"
        )