"""Add artifacts

Revision ID: 345551145801
Revises: 2bb529fb88b6
Create Date: 2025-08-05 05:48:29.768739

"""

import sqlalchemy as sa
from alembic import op

from ai.backend.manager.models.base import GUID

# revision identifiers, used by Alembic.
revision = "345551145801"
down_revision = "2bb529fb88b6"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "artifact_revisions",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("artifact_id", GUID(), nullable=False),
        sa.Column("version", sa.String(), nullable=False),
        sa.Column("size", sa.BigInteger(), nullable=True, default=None),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("readme", sa.TEXT(), nullable=True, default=None),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("readonly", sa.Boolean(), default=False, nullable=False),
        sa.UniqueConstraint("artifact_id", "version", name="uq_artifact_id_version"),
    )
    op.create_index(
        op.f("ix_artifact_revisions_artifact_id"),
        "artifact_revisions",
        ["artifact_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_artifact_revisions_created_at"), "artifact_revisions", ["created_at"], unique=False
    )
    op.create_index(
        op.f("ix_artifact_revisions_updated_at"), "artifact_revisions", ["updated_at"], unique=False
    )
    op.create_index(
        op.f("ix_artifact_revisions_status"), "artifact_revisions", ["status"], unique=False
    )

    op.create_table(
        "artifacts",
        sa.Column(
            "id",
            GUID(),
            server_default=sa.text("uuid_generate_v4()"),
            nullable=False,
        ),
        sa.Column(
            "type", sa.Enum("MODEL", "PACKAGE", "IMAGE", name="artifacttype"), nullable=False
        ),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("registry_id", GUID(), nullable=False),
        sa.Column("source_registry_id", GUID(), nullable=False),
        sa.Column("registry_type", sa.String(), nullable=False),
        sa.Column("source_registry_type", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=True),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_artifacts")),
    )
    op.create_index(op.f("ix_artifacts_name"), "artifacts", ["name"], unique=False)
    op.create_index(op.f("ix_artifacts_registry_id"), "artifacts", ["registry_id"], unique=False)
    op.create_index(
        op.f("ix_artifacts_source_registry_id"), "artifacts", ["source_registry_id"], unique=False
    )
    op.create_index(
        op.f("ix_artifacts_registry_type"), "artifacts", ["registry_type"], unique=False
    )
    op.create_index(
        op.f("ix_artifacts_source_registry_type"),
        "artifacts",
        ["source_registry_type"],
        unique=False,
    )
    op.create_index(op.f("ix_artifacts_type"), "artifacts", ["type"], unique=False)

    op.create_table(
        "association_artifacts_storages",
        sa.Column(
            "id",
            GUID(),
            server_default=sa.text("uuid_generate_v4()"),
            nullable=False,
        ),
        sa.Column("artifact_revision_id", GUID(), nullable=False),
        sa.Column("storage_id", GUID(), nullable=False),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_association_artifacts_storages")),
        sa.UniqueConstraint("artifact_revision_id", name="uq_artifact_revision_id"),
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("association_artifacts_storages")

    op.drop_index(op.f("ix_artifact_revisions_artifact_id"), table_name="artifact_revisions")
    op.drop_index(op.f("ix_artifact_revisions_created_at"), table_name="artifact_revisions")
    op.drop_index(op.f("ix_artifact_revisions_updated_at"), table_name="artifact_revisions")
    op.drop_index(op.f("ix_artifact_revisions_status"), table_name="artifact_revisions")
    op.drop_table("artifact_revisions")

    op.drop_index(op.f("ix_artifacts_type"), table_name="artifacts")
    op.drop_index(op.f("ix_artifacts_name"), table_name="artifacts")
    op.drop_index(op.f("ix_artifacts_registry_id"), table_name="artifacts")
    op.drop_index(op.f("ix_artifacts_source_registry_id"), table_name="artifacts")
    op.drop_index(op.f("ix_artifacts_registry_type"), table_name="artifacts")
    op.drop_index(op.f("ix_artifacts_source_registry_type"), table_name="artifacts")
    op.drop_table("artifacts")
    op.execute("DROP TYPE IF EXISTS artifacttype")
    # ### end Alembic commands ###
