"""Persist actor-scoped Model Lab execution identities."""

import sqlalchemy as sa

from alembic import op

revision = "20260924_0038"
down_revision = "20260923_0037"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.create_table(
        "lab_execution_runs",
        sa.Column("id", sa.Uuid(), nullable=False),
        sa.Column("actor_subject", sa.String(length=255), nullable=False),
        sa.Column("idempotency_key", sa.String(length=128), nullable=False),
        sa.Column("request_fingerprint", sa.String(length=64), nullable=False),
        sa.Column("status", sa.String(length=24), server_default="accepted", nullable=False),
        sa.Column("capability", sa.String(length=128), server_default="rag.embedding@1", nullable=False),
        sa.Column("target_snapshot", sa.JSON(), nullable=False),
        sa.Column("input_sha256", sa.String(length=64), nullable=False),
        sa.Column("input_count", sa.Integer(), nullable=False),
        sa.Column("input_encoded_bytes", sa.Integer(), nullable=False),
        sa.Column("payload_reference", sa.String(length=128), nullable=True),
        sa.Column("input_payload_expires_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("gateway_request_id", sa.Uuid(), nullable=False),
        sa.Column("timeout_seconds", sa.Integer(), server_default="60", nullable=False),
        sa.Column("queue_timeout_seconds", sa.Integer(), server_default="10", nullable=False),
        sa.Column("retention_seconds", sa.Integer(), server_default="300", nullable=False),
        sa.Column("output_metadata", sa.JSON(), nullable=True),
        sa.Column("worker_observations", sa.JSON(), nullable=True),
        sa.Column("failure_code", sa.String(length=64), nullable=True),
        sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("dispatch_started_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("heartbeat_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("cancellation_requested_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("cancellation_acknowledged_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
        sa.CheckConstraint(
            "status IN ('accepted', 'preparing', 'queued', 'running', 'cancel_requested', "
            "'completed', 'failed', 'cancelled', 'interrupted')",
            name="ck_lab_execution_run_status",
        ),
        sa.CheckConstraint("input_count BETWEEN 1 AND 8", name="ck_lab_execution_run_input_count"),
        sa.CheckConstraint("timeout_seconds BETWEEN 5 AND 300", name="ck_lab_execution_run_timeout"),
        sa.CheckConstraint("queue_timeout_seconds BETWEEN 1 AND 60", name="ck_lab_execution_run_queue_timeout"),
        sa.CheckConstraint("retention_seconds BETWEEN 60 AND 3600", name="ck_lab_execution_run_retention"),
        sa.ForeignKeyConstraint(
            ["gateway_request_id"], ["gateway_requests.id"], ondelete="RESTRICT"
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("actor_subject", "idempotency_key", name="uq_lab_run_actor_idempotency"),
        sa.UniqueConstraint("gateway_request_id", name="uq_lab_execution_run_gateway_request"),
    )
    op.create_index("ix_lab_execution_runs_actor_subject", "lab_execution_runs", ["actor_subject"])
    op.create_index("ix_lab_execution_runs_status", "lab_execution_runs", ["status"])
    op.create_index("ix_lab_execution_runs_gateway_request_id", "lab_execution_runs", ["gateway_request_id"])
    op.create_index("ix_lab_execution_runs_actor_status_created", "lab_execution_runs", ["actor_subject", "status", "created_at"])
    op.add_column(
        "serving_jobs",
        sa.Column("lab_execution_run_id", sa.Uuid(), nullable=True),
    )
    op.create_foreign_key(
        "fk_serving_jobs_lab_execution_run_id_lab_execution_runs",
        "serving_jobs",
        "lab_execution_runs",
        ["lab_execution_run_id"],
        ["id"],
        ondelete="SET NULL",
    )
    op.create_index(
        "ix_serving_jobs_lab_execution_run_id", "serving_jobs", ["lab_execution_run_id"]
    )
    if op.get_bind().dialect.name == "postgresql":
        op.execute(
            "REVOKE DELETE, TRUNCATE ON lab_execution_runs FROM modelforge_runtime"
        )
        op.execute(
            "GRANT SELECT, INSERT, UPDATE ON lab_execution_runs TO modelforge_runtime"
        )


def downgrade() -> None:
    bind = op.get_bind()
    if bind.scalar(sa.select(sa.func.count()).select_from(sa.table("lab_execution_runs", sa.column("id")))):
        raise RuntimeError(
            "LAB run history exists; restore the pre-migration backup instead of deleting it."
        )
    op.drop_index("ix_serving_jobs_lab_execution_run_id", table_name="serving_jobs")
    op.drop_constraint(
        "fk_serving_jobs_lab_execution_run_id_lab_execution_runs",
        "serving_jobs",
        type_="foreignkey",
    )
    op.drop_column("serving_jobs", "lab_execution_run_id")
    op.drop_index("ix_lab_execution_runs_actor_status_created", table_name="lab_execution_runs")
    op.drop_index("ix_lab_execution_runs_gateway_request_id", table_name="lab_execution_runs")
    op.drop_index("ix_lab_execution_runs_status", table_name="lab_execution_runs")
    op.drop_index("ix_lab_execution_runs_actor_subject", table_name="lab_execution_runs")
    op.drop_table("lab_execution_runs")