"""add claim fields to story/chapter image intents

Revision ID: 0010_intent_claims
Revises: 0009_si_constraints
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op


revision: str = "0010_intent_claims"
down_revision: Union[str, Sequence[str], None] = "0009_si_constraints"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _add_claim_columns(table_name: str) -> None:
    op.add_column(table_name, sa.Column("claim_token", sa.String(), nullable=True))
    op.add_column(
        table_name, sa.Column("claimed_at", sa.DateTime(timezone=True), nullable=True)
    )
    op.add_column(
        table_name,
        sa.Column(
            "attempt_count",
            sa.Integer(),
            nullable=True,
            server_default="0",
        ),
    )
    op.execute(f"UPDATE {table_name} SET attempt_count = 0 WHERE attempt_count IS NULL")
    op.alter_column(table_name, "attempt_count", nullable=False, server_default=None)


def _drop_claim_columns(table_name: str) -> None:
    op.drop_column(table_name, "attempt_count")
    op.drop_column(table_name, "claimed_at")
    op.drop_column(table_name, "claim_token")


def upgrade() -> None:
    _add_claim_columns("story_image_intents")
    _add_claim_columns("chapter_cover_intents")


def downgrade() -> None:
    _drop_claim_columns("chapter_cover_intents")
    _drop_claim_columns("story_image_intents")