app/generation/models.py

from __future__ import annotations

from datetime import datetime, timezone
from uuid import uuid4

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

from app.security.models import Base


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


def new_id() -> str:
    return str(uuid4())


class GenerationRequest(Base):
    """Immutable validated intent, scoped to one authoritative workspace."""

    __tablename__ = "generation_requests"
    __table_args__ = (
        UniqueConstraint(
            "workspace_id", "idempotency_key", name="uq_generation_request_workspace_idempotency"
        ),
        Index("ix_generation_requests_workspace_created", "workspace_id", "created_at"),
        Index("ix_generation_requests_workspace_status", "workspace_id", "status"),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
    workspace_id: Mapped[str] = mapped_column(
        String(36), ForeignKey("workspaces.id", ondelete="RESTRICT"), nullable=False
    )
    created_by_user_id: Mapped[str] = mapped_column(
        String(36), ForeignKey("users.id", ondelete="RESTRICT"), nullable=False
    )
    provider: Mapped[str] = mapped_column(String(64), nullable=False)
    model_id: Mapped[str] = mapped_column(String(255), nullable=False)
    modality: Mapped[str] = mapped_column(String(32), nullable=False)
    input_asset_id: Mapped[str | None] = mapped_column(
        String(36), ForeignKey("media_assets.id", ondelete="RESTRICT")
    )
    project_id: Mapped[str | None] = mapped_column(
        String(36), ForeignKey("projects.id", ondelete="RESTRICT")
    )
    product_surface: Mapped[str] = mapped_column(String(32), nullable=False, default="generation")
    # This JSON contains only adapter-validated, non-secret input.  Its
    # request fingerprint is authoritative for idempotency conflict checks.
    spec_json: Mapped[dict[str, object]] = mapped_column("spec", JSON, nullable=False, default=dict)
    request_fingerprint: Mapped[str] = mapped_column(String(64), nullable=False)
    idempotency_key: Mapped[str] = mapped_column(String(255), nullable=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="queued")
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=utcnow
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=utcnow, onupdate=utcnow
    )
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))


class GenerationJob(Base):
    """One durable execution for one request; retries keep this identity."""

    __tablename__ = "generation_jobs"
    __table_args__ = (
        UniqueConstraint("generation_request_id", name="uq_generation_job_request"),
        # An external worker job is an opaque provider identity, not a
        # workspace-scoped client value. Binding it once prevents a worker
        # status/output from being attached to another tenant's logical job.
        UniqueConstraint("provider", "external_job_id", name="uq_generation_job_provider_external"),
        Index("ix_generation_jobs_workspace_status", "workspace_id", "status"),
        Index("ix_generation_jobs_next_attempt", "status", "next_attempt_at"),
        Index("ix_generation_jobs_external", "provider", "external_job_id"),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
    generation_request_id: Mapped[str] = mapped_column(
        String(36),
        ForeignKey("generation_requests.id", ondelete="CASCADE"),
        nullable=False,
    )
    workspace_id: Mapped[str] = mapped_column(
        String(36), ForeignKey("workspaces.id", ondelete="RESTRICT"), nullable=False
    )
    provider: Mapped[str] = mapped_column(String(64), nullable=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="queued")
    attempt_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    max_attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=3)
    next_attempt_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    external_job_id: Mapped[str | None] = mapped_column(String(255))
    # Safe metadata only.  Bearer-like upload URLs or credentials belong in a
    # future secret store, never in this row or externally serialised views.
    provider_metadata_json: Mapped[dict[str, object]] = mapped_column(
        "provider_metadata", JSON, nullable=False, default=dict
    )
    output_asset_id: Mapped[str | None] = mapped_column(
        String(36), ForeignKey("media_assets.id", ondelete="RESTRICT")
    )
    error_code: Mapped[str | None] = mapped_column(String(100))
    error_message: Mapped[str | None] = mapped_column(Text)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=utcnow
    )
    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=utcnow, onupdate=utcnow
    )


class GenerationJobAttempt(Base):
    """Auditable retry history without creating new logical jobs."""

    __tablename__ = "generation_job_attempts"
    __table_args__ = (
        UniqueConstraint(
            "generation_job_id", "attempt_number", name="uq_generation_job_attempt_number"
        ),
        Index("ix_generation_job_attempts_job", "generation_job_id", "attempt_number"),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
    generation_job_id: Mapped[str] = mapped_column(
        String(36), ForeignKey("generation_jobs.id", ondelete="CASCADE"), nullable=False
    )
    attempt_number: Mapped[int] = mapped_column(Integer, nullable=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False)
    brand_kit_version_id: Mapped[str | None] = mapped_column(String(36), nullable=True)
    external_job_id: Mapped[str | None] = mapped_column(String(128))
    error_message: Mapped[str | None] = mapped_column(Text)
    provider_request_id: Mapped[str | None] = mapped_column(String(255))
    started_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=utcnow
    )
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))