app/models/database.py

"""
Ukweli — SQLAlchemy ORM Models
Defines the relational schema for PostgreSQL per Architecture Section 3.3.4.
"""

from __future__ import annotations

import uuid
from datetime import datetime

from sqlalchemy import (
    Boolean,
    DateTime,
    Enum,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    func,
)
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class Base(DeclarativeBase):
    """Declarative base for all ORM models."""

    pass


class Document(Base):
    """
    Catalog of all ingested documents.
    Maps to the Metadata Registry described in Section 3.3.4.
    """

    __tablename__ = "documents"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    title: Mapped[str] = mapped_column(String(1024), nullable=False)
    source_url: Mapped[str | None] = mapped_column(String(2048), nullable=True)
    pdf_url: Mapped[str | None] = mapped_column(String(2048), nullable=True)
    fiscal_year: Mapped[str | None] = mapped_column(String(20), nullable=True, index=True)
    auditee: Mapped[str | None] = mapped_column(String(512), nullable=True, index=True)
    report_type: Mapped[str] = mapped_column(
        Enum(
            "financial_audit",
            "performance_audit",
            "special_audit",
            "public_debt_audit",
            "county_audit",
            "summary_report",
            "strategic_document",
            "annual_corporate_report",
            name="report_type_enum",
        ),
        nullable=False,
        default="financial_audit",
    )
    page_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
    published_date: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    checksum_sha256: Mapped[str | None] = mapped_column(String(64), nullable=True, unique=True)
    status: Mapped[str] = mapped_column(
        Enum(
            "pending",
            "downloading",
            "parsing",
            "parsed",
            "embedding",
            "ready",
            "failed",
            name="document_status_enum",
        ),
        nullable=False,
        default="pending",
        index=True,
    )
    raw_storage_path: Mapped[str | None] = mapped_column(String(1024), nullable=True)
    language: Mapped[str] = mapped_column(String(5), nullable=False, default="en")
    file_size_bytes: Mapped[int | None] = mapped_column(Integer, nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False
    )

    # Relationships
    chunks: Mapped[list[Chunk]] = relationship("Chunk", back_populates="document", cascade="all, delete-orphan")
    ingestion_logs: Mapped[list[IngestionLog]] = relationship(
        "IngestionLog", back_populates="document", cascade="all, delete-orphan"
    )

    __table_args__ = (
        Index("ix_documents_auditee_fy", "auditee", "fiscal_year"),
    )


class Chunk(Base):
    """
    Individual content chunks derived from parsed documents.
    Schema matches Section 3.2.4 Chunk Metadata Schema.
    """

    __tablename__ = "chunks"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    document_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False, index=True
    )
    chunk_type: Mapped[str] = mapped_column(
        Enum(
            "narrative",
            "table",
            "finding",
            "recommendation",
            "legal_reference",
            name="chunk_type_enum",
        ),
        nullable=False,
    )
    content: Mapped[str] = mapped_column(Text, nullable=False)
    page_range_start: Mapped[int | None] = mapped_column(Integer, nullable=True)
    page_range_end: Mapped[int | None] = mapped_column(Integer, nullable=True)
    section_heading: Mapped[str | None] = mapped_column(String(512), nullable=True)
    entities: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
    language: Mapped[str] = mapped_column(String(5), nullable=False, default="en")
    audit_period: Mapped[str | None] = mapped_column(String(20), nullable=True)
    auditee: Mapped[str | None] = mapped_column(String(512), nullable=True)
    finding_category: Mapped[str | None] = mapped_column(String(64), nullable=True)
    embedding_vector_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
    token_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
    chunk_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )

    # Relationships
    document: Mapped[Document] = relationship("Document", back_populates="chunks")

    __table_args__ = (
        Index("ix_chunks_document_index", "document_id", "chunk_index"),
        Index("ix_chunks_auditee_period", "auditee", "audit_period"),
        # Index(
        #     "ix_chunks_content_fts",
        #     func.to_tsvector("english", "content"),
        #     postgresql_using="gin",
        # ),
    )


class IngestionLog(Base):
    """
    Audit trail for every document processing step.
    Provides transparency per Section 3.3.4.
    """

    __tablename__ = "ingestion_logs"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    document_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False, index=True
    )
    stage: Mapped[str] = mapped_column(
        String(64), nullable=False
    )  # e.g. "download", "parse", "chunk", "embed"
    status: Mapped[str] = mapped_column(String(32), nullable=False)  # "started", "completed", "failed"
    started_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    worker_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
    metadata_extra: Mapped[dict | None] = mapped_column(JSONB, nullable=True)

    # Relationships
    document: Mapped[Document] = relationship("Document", back_populates="ingestion_logs")


class QueryLog(Base):
    """
    Immutable audit trail of every RAG query.
    Required by Section 7.2 — transparency logging (WORM-style).
    """

    __tablename__ = "query_logs"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    query_text: Mapped[str] = mapped_column(Text, nullable=False)
    language: Mapped[str] = mapped_column(String(5), nullable=False, default="en")
    mode: Mapped[str] = mapped_column(String(20), nullable=False, default="concise")
    user_tier: Mapped[str] = mapped_column(String(20), nullable=False, default="public")
    user_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True), ForeignKey("users.id"), nullable=True, index=True
    )
    api_key_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True), ForeignKey("api_keys.id"), nullable=True
    )
    fingerprint: Mapped[str | None] = mapped_column(String(16), nullable=True, index=True)
    filters: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
    answer: Mapped[str | None] = mapped_column(Text, nullable=True)
    citations: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
    chunks_considered: Mapped[int | None] = mapped_column(Integer, nullable=True)
    latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
    confidence: Mapped[str | None] = mapped_column(String(32), nullable=True)
    llm_model_used: Mapped[str | None] = mapped_column(String(128), nullable=True)
    llm_tokens_used: Mapped[int | None] = mapped_column(Integer, nullable=True)
    conversation_id: Mapped[uuid.UUID | None] = mapped_column(UUID(as_uuid=True), nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )

    # Relationships
    feedbacks: Mapped[list[Feedback]] = relationship("Feedback", back_populates="query_log")


class Feedback(Base):
    """
    User feedback and corrections on RAG responses.
    Maps to Section 4.2 POST /feedback.
    """

    __tablename__ = "feedbacks"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    query_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("query_logs.id", ondelete="CASCADE"), nullable=False, index=True
    )
    rating: Mapped[int] = mapped_column(Integer, nullable=False)  # -1, 0, 1
    correction: Mapped[str | None] = mapped_column(Text, nullable=True)
    suggested_citation: Mapped[str | None] = mapped_column(Text, nullable=True)
    review_status: Mapped[str] = mapped_column(
        String(20), nullable=False, default="pending"
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )

    # Relationships
    query_log: Mapped[QueryLog] = relationship("QueryLog", back_populates="feedbacks")


class PageHash(Base):
    """
    Stores content hashes of crawled pages for change detection.
    Used by the change_detection service.
    """

    __tablename__ = "page_hashes"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    url_path: Mapped[str] = mapped_column(String(2048), nullable=False, unique=True, index=True)
    content_hash: Mapped[str] = mapped_column(String(64), nullable=False)
    last_checked_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )
    changed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)