"""RAG models: sources and chunks with pgvector embeddings."""

from __future__ import annotations

from datetime import datetime
from typing import Any, Optional

from sqlalchemy import (
    DateTime,
    ForeignKey,
    Integer,
    JSON,
    String,
    Text,
    UniqueConstraint,
    func,
)
from sqlalchemy.orm import Mapped, mapped_column, relationship

from app.core.database import Base

try:
    from pgvector.sqlalchemy import Vector

    EMBEDDING_DIM = 1536
    _VectorType = Vector(EMBEDDING_DIM)
except Exception:  # pragma: no cover
    EMBEDDING_DIM = 1536
    _VectorType = JSON  # type: ignore[misc, assignment]


class RagSource(Base):
    __tablename__ = "rag_sources"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    scope: Mapped[str] = mapped_column(String(16), nullable=False, index=True)  # global | project
    project_id: Mapped[Optional[int]] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=True, index=True
    )
    kind: Mapped[str] = mapped_column(String(16), nullable=False)  # upload | url
    title: Mapped[str] = mapped_column(String(512), default="")
    uri: Mapped[Optional[str]] = mapped_column(String(2048), nullable=True)
    path: Mapped[Optional[str]] = mapped_column(String(2048), nullable=True)
    gost_id: Mapped[Optional[str]] = mapped_column(String(128), nullable=True, index=True)
    status: Mapped[str] = mapped_column(String(32), default="pending")  # pending|indexing|ready|error
    error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    meta: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )

    chunks: Mapped[list[RagChunk]] = relationship(
        back_populates="source", cascade="all, delete-orphan"
    )


class RagChunk(Base):
    __tablename__ = "rag_chunks"
    __table_args__ = (UniqueConstraint("source_id", "chunk_index", name="uq_rag_chunk_source_idx"),)

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    source_id: Mapped[int] = mapped_column(
        ForeignKey("rag_sources.id", ondelete="CASCADE"), index=True, nullable=False
    )
    chunk_index: Mapped[int] = mapped_column(Integer, default=0)
    text: Mapped[str] = mapped_column(Text, default="")
    embedding: Mapped[Optional[Any]] = mapped_column(_VectorType, nullable=True)
    meta: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)

    source: Mapped[RagSource] = relationship(back_populates="chunks")