"""add repository data key epochs

Revision ID: 0007_repository_data_key_epochs
Revises: 0006_local_sources_only
"""

import sqlalchemy as sa
from alembic import op

revision = "0007_repository_data_key_epochs"
down_revision = "0006_local_sources_only"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column("repositories", sa.Column("active_data_key_id", sa.String(36), nullable=True))
    op.add_column("backups", sa.Column("data_key_id", sa.String(36), nullable=True))
    op.create_table(
        "repository_data_key_epochs",
        sa.Column("id", sa.String(36), primary_key=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("CURRENT_TIMESTAMP"),
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("CURRENT_TIMESTAMP"),
        ),
        sa.Column(
            "repository_id",
            sa.String(36),
            sa.ForeignKey("repositories.id", ondelete="RESTRICT"),
            nullable=False,
        ),
        sa.Column("key_id", sa.String(36), nullable=False),
        sa.Column("state", sa.String(16), nullable=False),
        sa.Column("retired_at", sa.DateTime(timezone=True)),
        sa.UniqueConstraint(
            "repository_id", "key_id", name="uq_repository_data_key_epochs_repository_key_epoch"
        ),
        sa.CheckConstraint(
            "state IN ('active','retired')",
            name="ck_repository_data_key_epochs_repository_data_key_epoch_state",
        ),
    )
    op.create_index(
        "ix_repository_data_key_epochs_repository_id",
        "repository_data_key_epochs",
        ["repository_id"],
    )
    op.create_index(
        "uq_repository_data_key_epochs_active",
        "repository_data_key_epochs",
        ["repository_id"],
        unique=True,
        sqlite_where=sa.text("state = 'active'"),
    )


def downgrade() -> None:
    connection = op.get_bind()
    epochs = sa.table("repository_data_key_epochs")
    repositories = sa.table("repositories", sa.column("active_data_key_id"))
    backups = sa.table("backups", sa.column("data_key_id"))
    epoch_count = connection.scalar(sa.select(sa.func.count()).select_from(epochs))
    active_key_count = connection.scalar(
        sa.select(sa.func.count())
        .select_from(repositories)
        .where(repositories.c.active_data_key_id.is_not(None))
    )
    backup_key_count = connection.scalar(
        sa.select(sa.func.count()).select_from(backups).where(backups.c.data_key_id.is_not(None))
    )
    if epoch_count or active_key_count or backup_key_count:
        raise RuntimeError("cannot downgrade while repository data key metadata exists")
    op.drop_index("uq_repository_data_key_epochs_active", table_name="repository_data_key_epochs")
    op.drop_index(
        "ix_repository_data_key_epochs_repository_id", table_name="repository_data_key_epochs"
    )
    op.drop_table("repository_data_key_epochs")
    op.drop_column("backups", "data_key_id")
    op.drop_column("repositories", "active_data_key_id")