backend/database.py

"""
Antigravity Notebook - Database Models
SQLAlchemy models for PostgreSQL database.
"""

from sqlalchemy import create_engine, Column, String, Integer, Text, ForeignKey, DateTime, CheckConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.dialects.postgresql import UUID, JSONB
from datetime import datetime
import uuid

from backend.config import settings

# Database engine
engine = create_engine(settings.DATABASE_URL, echo=True)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()


class Notebook(Base):
    """Notebook model - container for a project"""
    __tablename__ = "notebooks"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(String(255), nullable=False)
    description = Column(Text, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Relationships
    sources = relationship("Source", back_populates="notebook", cascade="all, delete-orphan")
    chat_messages = relationship("ChatMessage", back_populates="notebook", cascade="all, delete-orphan")

    def __repr__(self):
        return f"<Notebook(id={self.id}, name={self.name})>"


class Source(Base):
    """Source model - a file (PDF), URL, or text belonging to a notebook"""
    __tablename__ = "sources"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    notebook_id = Column(UUID(as_uuid=True), ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False)
    source_type = Column(String(10), nullable=False)
    filename = Column(String(255), nullable=True)
    url = Column(Text, nullable=True)
    content_hash = Column(String(64), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    source_metadata = Column(JSONB, default={})

    __table_args__ = (
        CheckConstraint("source_type IN ('pdf', 'url', 'text')", name="check_source_type"),
    )

    # Relationships
    notebook = relationship("Notebook", back_populates="sources")
    latent_tensors = relationship("LatentTensor", back_populates="source", cascade="all, delete-orphan")

    def __repr__(self):
        return f"<Source(id={self.id}, type={self.source_type}, filename={self.filename})>"


class LatentTensor(Base):
    """LatentTensor model - CLaRa-compressed representation of a segment"""
    __tablename__ = "latent_tensors"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    source_id = Column(UUID(as_uuid=True), ForeignKey("sources.id", ondelete="CASCADE"), nullable=False)
    tensor_path = Column(String(512), nullable=False)
    segment_index = Column(Integer, nullable=False)
    token_count = Column(Integer, nullable=False)
    original_text_length = Column(Integer, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    tensor_metadata = Column(JSONB, default={})

    # Relationships
    source = relationship("Source", back_populates="latent_tensors")

    def __repr__(self):
        return f"<LatentTensor(id={self.id}, source_id={self.source_id}, tokens={self.token_count})>"


class ChatMessage(Base):
    """ChatMessage model - stores chat history"""
    __tablename__ = "chat_messages"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    notebook_id = Column(UUID(as_uuid=True), ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False)
    role = Column(String(20), nullable=False)
    content = Column(Text, nullable=False)
    sources_used = Column(JSONB, default=[])
    created_at = Column(DateTime, default=datetime.utcnow)

    __table_args__ = (
        CheckConstraint("role IN ('user', 'assistant')", name="check_role"),
    )

    # Relationships
    notebook = relationship("Notebook", back_populates="chat_messages")

    def __repr__(self):
        return f"<ChatMessage(id={self.id}, role={self.role})>"


def get_db():
    """Dependency for FastAPI to get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db():
    """Initialize database - create all tables"""
    Base.metadata.create_all(bind=engine)
    print("✅ Database initialized successfully!")


if __name__ == "__main__":
    init_db()