"""Add artifact unique constraints [c22561cbb3a9].

Revision ID: c22561cbb3a9
Revises: 0.68.1
Create Date: 2024-10-17 16:41:25.053677

"""

from collections import defaultdict
from typing import Dict, Set

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "c22561cbb3a9"
down_revision = "0.68.1"
branch_labels = None
depends_on = None


def resolve_duplicate_versions() -> None:
    """Resolve duplicate artifact versions."""
    connection = op.get_bind()

    meta = sa.MetaData()
    meta.reflect(
        bind=op.get_bind(),
        only=("artifact_version",),
    )

    artifact_version_table = sa.Table("artifact_version", meta)
    query = sa.select(
        artifact_version_table.c.id,
        artifact_version_table.c.artifact_id,
        artifact_version_table.c.version,
    )

    versions_per_artifact: Dict[str, Set[str]] = defaultdict(set)

    for id, artifact_id, version in connection.execute(query).fetchall():
        versions = versions_per_artifact[artifact_id]
        if version in versions:
            for suffix_length in range(4, len(id)):
                new_version = f"{version}-{id[:suffix_length]}"
                if new_version not in versions:
                    version = new_version
                    break

            connection.execute(
                sa.update(artifact_version_table)
                .where(artifact_version_table.c.id == id)
                .values(version=version)
            )

        versions.add(version)


def upgrade() -> None:
    """Upgrade database schema and/or data, creating a new revision."""
    # ### commands auto generated by Alembic - please adjust! ###
    resolve_duplicate_versions()

    with op.batch_alter_table("artifact", schema=None) as batch_op:
        batch_op.create_unique_constraint("unique_artifact_name", ["name"])

    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
        batch_op.create_unique_constraint(
            "unique_version_for_artifact_id", ["version", "artifact_id"]
        )

    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade database schema and/or data back to the previous revision."""
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
        batch_op.drop_constraint(
            "unique_version_for_artifact_id", type_="unique"
        )

    with op.batch_alter_table("artifact", schema=None) as batch_op:
        batch_op.drop_constraint("unique_artifact_name", type_="unique")

    # ### end Alembic commands ###
