imagerepo/app/models.py

from pgvector.sqlalchemy import Vector

from . import db

# tag to object relationship table
image_tags = db.Table(
    "image_tags",
    db.Column(
        "image_id", db.Integer,
        db.ForeignKey("image_record.id", ondelete="CASCADE"),
        primary_key=True,
    ),
    db.Column(
        "tag_id", db.Integer,
        db.ForeignKey("tag.id", ondelete="CASCADE"),
        primary_key=True,
    ),
    db.UniqueConstraint("image_id", "tag_id", name="uq_image_tag"),
)

class ImageRecord(db.Model):
    __tablename__ = "image_record"
    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(255), nullable=False)
    filepath = db.Column(db.String(512), nullable=False)
    thumb_path = db.Column(db.String(512))
    hash = db.Column(db.String(64), unique=True, nullable=False)  # SHA256
    perceptual_hash = db.Column(db.String(64), nullable=True)  # hex of 256-bit hash
    file_size = db.Column(db.Integer, nullable=True)
    width = db.Column(db.Integer, nullable=True)
    height = db.Column(db.Integer, nullable=True)
    format = db.Column(db.String(10), nullable=True)
    camera_model = db.Column(db.String(255), nullable=True)
    taken_at = db.Column(db.DateTime, nullable=True)
    imported_at = db.Column(db.DateTime, default=db.func.now())
    is_thumbnail = db.Column(db.Boolean, default=False)
    tags = db.relationship(
        "Tag",
        secondary=image_tags,
        back_populates="images",
        passive_deletes=True,
    )

class Tag(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(255), nullable=False)
    kind = db.Column(db.String(64), nullable=True)
    fandom_id = db.Column(db.Integer, db.ForeignKey("tag.id", ondelete="SET NULL"), nullable=True)
    fandom = db.relationship("Tag", remote_side="Tag.id", foreign_keys=[fandom_id])
    images = db.relationship(
        "ImageRecord",
        secondary=image_tags,
        back_populates="tags",
        passive_deletes=True,
    )

    __table_args__ = (
        db.Index(
            'tag_character_with_fandom_uniq', 'name', 'fandom_id',
            unique=True,
            postgresql_where=db.text("kind = 'character' AND fandom_id IS NOT NULL"),
        ),
        db.Index(
            'tag_character_null_fandom_uniq', 'name',
            unique=True,
            postgresql_where=db.text("kind = 'character' AND fandom_id IS NULL"),
        ),
        db.Index(
            'tag_other_kinds_uniq', 'name', 'kind',
            unique=True,
            postgresql_where=db.text("kind != 'character'"),
        ),
        db.Index('tag_kind_idx', 'kind'),
    )

    @property
    def display_name(self) -> str:
        """Human-readable rendering of this tag.

        After the bare-name refactor, `name` holds the display text directly for
        every kind except character+fandom, which gets the '(Fandom)' suffix
        appended here. Pre-migration data (with 'kind:' prefix in name) still
        renders sanely via the else branch — callers see the raw stored string.
        """
        if self.kind == 'character' and self.fandom is not None:
            return f"{self.name} ({self.fandom.name})"
        return self.name

class ArchiveRecord(db.Model):
    """
    Tracks archives we've processed so we can skip re-importing the same large file.
    Also holds a Tag that marks images that came from this archive.
    """
    __tablename__ = "archive_record"
    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(512), nullable=False)
    file_size = db.Column(db.BigInteger, nullable=False)
    hash = db.Column(db.String(64), unique=True, nullable=False)  # SHA256 of the archive file
    imported_at = db.Column(db.DateTime, default=db.func.now())
    artist = db.Column(db.String(255), nullable=True)
    tag_id = db.Column(db.Integer, db.ForeignKey("tag.id", ondelete="SET NULL"), nullable=True)
    tag = db.relationship("Tag")


class AppSettings(db.Model):
    """
    Key-value store for application settings and version tracking.
    Used to persist configuration and trigger data migrations on version changes.
    """
    __tablename__ = "app_settings"
    id = db.Column(db.Integer, primary_key=True)
    key = db.Column(db.String(64), unique=True, nullable=False)
    value = db.Column(db.Text, nullable=True)
    updated_at = db.Column(db.DateTime, default=db.func.now(), onupdate=db.func.now())


class ImportBatch(db.Model):
    """
    Groups related import tasks for batch tracking and reporting.
    One batch per scan operation.
    """
    __tablename__ = "import_batch"

    id = db.Column(db.Integer, primary_key=True)
    source_directory = db.Column(db.String(1024), nullable=False)
    status = db.Column(db.String(32), default='running', index=True)  # running, complete, failed, cancelled

    # Statistics
    total_files = db.Column(db.Integer, default=0)
    processed_files = db.Column(db.Integer, default=0)
    imported_count = db.Column(db.Integer, default=0)
    skipped_count = db.Column(db.Integer, default=0)
    failed_count = db.Column(db.Integer, default=0)

    # Timestamps
    started_at = db.Column(db.DateTime, default=db.func.now())
    completed_at = db.Column(db.DateTime, nullable=True)


class PostMetadata(db.Model):
    """
    Stores metadata about posts from external platforms (Patreon, SubscribeStar, HentaiFoundry).
    Links to a Tag with kind='post' (format: post:platform:artist:id).
    """
    __tablename__ = "post_metadata"

    id = db.Column(db.Integer, primary_key=True)

    # Link to the post tag (post:platform:artist:id)
    tag_id = db.Column(db.Integer, db.ForeignKey("tag.id", ondelete="CASCADE"), nullable=False, unique=True)
    tag = db.relationship("Tag", backref=db.backref("post_metadata", uselist=False))

    # Platform identification
    platform = db.Column(db.String(64), nullable=False)  # patreon, subscribestar, hentaifoundry
    post_id = db.Column(db.String(64), nullable=False)  # Original post ID on platform
    artist = db.Column(db.String(255), nullable=True)

    # Content
    title = db.Column(db.String(512), nullable=True)
    description = db.Column(db.Text, nullable=True)  # May contain HTML
    source_url = db.Column(db.String(1024), nullable=True)  # Link to original post

    # Statistics
    attachment_count = db.Column(db.Integer, default=0)  # Number of images in post

    # Timestamps
    published_at = db.Column(db.DateTime, nullable=True)  # Original publish date
    created_at = db.Column(db.DateTime, default=db.func.now())
    updated_at = db.Column(db.DateTime, default=db.func.now(), onupdate=db.func.now())


class SeriesPage(db.Model):
    """
    Links images to series tags with page numbers for ordered reading.
    An image can only belong to one series at a time.
    """
    __tablename__ = "series_page"

    id = db.Column(db.Integer, primary_key=True)

    # Link to the series tag (must be kind='series')
    series_tag_id = db.Column(db.Integer, db.ForeignKey("tag.id", ondelete="CASCADE"), nullable=False)
    series_tag = db.relationship("Tag", backref=db.backref("series_pages", lazy="dynamic"))

    # Link to the image (unique - image can only be in one series)
    image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, unique=True)
    image = db.relationship("ImageRecord", backref=db.backref("series_page", uselist=False))

    # Page number within the series
    page_number = db.Column(db.Integer, nullable=False)

    # Timestamps
    created_at = db.Column(db.DateTime, default=db.func.now())
    updated_at = db.Column(db.DateTime, default=db.func.now(), onupdate=db.func.now())

    __table_args__ = (
        # Each page number must be unique within a series
        db.UniqueConstraint('series_tag_id', 'page_number', name='uq_series_page_number'),
        db.Index('ix_series_page_series_tag', 'series_tag_id'),
    )


class ImportTask(db.Model):
    """
    Tracks the state of each file being processed through the import pipeline.
    Enables resume capability and parallel processing.
    """
    __tablename__ = "import_task"

    id = db.Column(db.Integer, primary_key=True)

    # File identification
    source_path = db.Column(db.String(1024), nullable=False, index=True)
    file_hash = db.Column(db.String(64), nullable=True, index=True)  # SHA256, computed during processing
    file_size = db.Column(db.BigInteger, nullable=True)

    # Classification
    # Values: 'scan', 'import_image', 'import_archive', 'thumbnail', 'sidecar'
    task_type = db.Column(db.String(32), nullable=False, index=True)

    # State tracking
    # Values: 'pending', 'queued', 'processing', 'complete', 'failed', 'skipped'
    status = db.Column(db.String(32), nullable=False, default='pending', index=True)

    # Celery task tracking
    celery_task_id = db.Column(db.String(64), nullable=True, index=True)

    # Context/metadata (JSON)
    # For images: {"artist": "...", "dest_dir": "..."}
    # For archives: {"artist": "...", "dest_dir": "...", "archive_tag": "..."}
    # For sidecars: {"image_id": 123, "sidecar_path": "..."}
    context = db.Column(db.JSON, nullable=True)

    # Batch relationship
    batch_id = db.Column(db.Integer, db.ForeignKey('import_batch.id', ondelete='SET NULL'), nullable=True)
    batch = db.relationship('ImportBatch', backref=db.backref('tasks', lazy='dynamic'))

    # Result tracking
    result_image_id = db.Column(db.Integer, db.ForeignKey('image_record.id', ondelete='SET NULL'), nullable=True)
    result_image = db.relationship('ImageRecord', foreign_keys=[result_image_id])
    error_message = db.Column(db.Text, nullable=True)
    retry_count = db.Column(db.Integer, default=0)

    # Timestamps
    created_at = db.Column(db.DateTime, default=db.func.now(), index=True)
    started_at = db.Column(db.DateTime, nullable=True)
    completed_at = db.Column(db.DateTime, nullable=True)

    # Unique constraint: prevent duplicate tasks for same file+type combination
    # Note: file_hash can be NULL initially (computed during processing)
    __table_args__ = (
        db.Index('ix_import_task_status_type', 'status', 'task_type'),
        db.Index('ix_import_task_batch_status', 'batch_id', 'status'),
    )


class ImageTagPrediction(db.Model):
    __tablename__ = "image_tag_prediction"
    id = db.Column(db.Integer, primary_key=True)
    image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False)
    tag_name = db.Column(db.Text, nullable=False)
    tag_category = db.Column(db.Text, nullable=False)  # general/character/copyright/rating/meta
    confidence = db.Column(db.Float, nullable=False)
    model_version = db.Column(db.Text, nullable=False)
    created_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)

    __table_args__ = (
        db.Index('idx_tag_predictions_image', 'image_id'),
        db.Index('idx_tag_predictions_tag', 'tag_name'),
        db.Index('idx_tag_predictions_model', 'model_version'),
    )


class ImageEmbedding(db.Model):
    __tablename__ = "image_embedding"
    image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), primary_key=True)
    model_version = db.Column(db.Text, primary_key=True)
    embedding = db.Column(Vector(1152), nullable=False)
    created_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)


class TagReferenceEmbedding(db.Model):
    __tablename__ = "tag_reference_embedding"
    tag_name = db.Column(db.Text, primary_key=True)
    model_version = db.Column(db.Text, primary_key=True)
    tag_kind = db.Column(db.Text, nullable=True)
    centroid = db.Column(Vector(1152), nullable=False)
    reference_count = db.Column(db.Integer, nullable=False)
    computed_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)


class SuggestionFeedback(db.Model):
    __tablename__ = "suggestion_feedback"
    id = db.Column(db.Integer, primary_key=True)
    image_id = db.Column(db.Integer, db.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False)
    tag_name = db.Column(db.Text, nullable=False)
    suggestion_source = db.Column(db.Text, nullable=False)  # wd14 | embedding_similarity
    confidence = db.Column(db.Float, nullable=False)
    decision = db.Column(db.Text, nullable=False)  # accepted | rejected
    decided_at = db.Column(db.DateTime(timezone=True), server_default=db.func.now(), nullable=False)

    __table_args__ = (
        db.Index('idx_feedback_image', 'image_id'),
        db.Index('idx_feedback_tag', 'tag_name'),
    )


class TagSuggestionConfig(db.Model):
    __tablename__ = "tag_suggestion_config"
    key = db.Column(db.Text, primary_key=True)
    value = db.Column(db.Text, nullable=False)
    description = db.Column(db.Text, nullable=True)