"""add resize metadata to instance disk addons

Revision ID: 024_instance_disk_addon_resize
Revises: 023_instance_disk_addons
Create Date: 2026-06-13
"""

from alembic import op
import sqlalchemy as sa

revision = "024_instance_disk_addon_resize"
down_revision = "023_instance_disk_addons"
branch_labels = None
depends_on = None


def _columns(table_name: str) -> set[str]:
    bind = op.get_bind()
    return {column["name"] for column in sa.inspect(bind).get_columns(table_name)}


def upgrade() -> None:
    bind = op.get_bind()
    tables = set(sa.inspect(bind).get_table_names())
    if "instance_disk_addons" not in tables:
        return

    columns = _columns("instance_disk_addons")
    if "kind" not in columns:
        op.add_column(
            "instance_disk_addons",
            sa.Column("kind", sa.String(24), nullable=False, server_default="create_disk"),
        )
    if "source_size_gb" not in columns:
        op.add_column("instance_disk_addons", sa.Column("source_size_gb", sa.Integer(), nullable=True))
    if "target_size_gb" not in columns:
        op.add_column("instance_disk_addons", sa.Column("target_size_gb", sa.Integer(), nullable=True))
    if "delta_size_gb" not in columns:
        op.add_column("instance_disk_addons", sa.Column("delta_size_gb", sa.Integer(), nullable=True))

    op.execute(
        """
        UPDATE instance_disk_addons
        SET
          source_size_gb = COALESCE(source_size_gb, 0),
          target_size_gb = COALESCE(target_size_gb, size_gb),
          delta_size_gb = COALESCE(delta_size_gb, size_gb)
        WHERE kind = 'create_disk'
        """
    )


def downgrade() -> None:
    bind = op.get_bind()
    tables = set(sa.inspect(bind).get_table_names())
    if "instance_disk_addons" not in tables:
        return

    columns = _columns("instance_disk_addons")
    if "delta_size_gb" in columns:
        op.drop_column("instance_disk_addons", "delta_size_gb")
    if "target_size_gb" in columns:
        op.drop_column("instance_disk_addons", "target_size_gb")
    if "source_size_gb" in columns:
        op.drop_column("instance_disk_addons", "source_size_gb")
    if "kind" in columns:
        op.drop_column("instance_disk_addons", "kind")
