"""switch to timezone aware

Revision ID: cb3d6f0cdb86
Revises: 54b31e87fe33
Create Date: 2024-12-06 17:34:51.630311

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'cb3d6f0cdb86'
down_revision = '54b31e87fe33'
branch_labels = None
depends_on = None


def upgrade():
    def alter_column_to_timezone_aware(table_name, column_name):
        with op.batch_alter_table(table_name, schema=None) as batch_op:
            batch_op.alter_column(
                column_name,
                type_=sa.DateTime(timezone=True)
            )

    # AbstractConfiguration derived tables
    configuration_tables = [
        "automation",
        "cloud_account",
        "bridge_conf",
        "country",
        "group",
        "mirror_list",
        "onion",
        "origin",
        "pool",
        "static_origin",
        "webhook"
    ]
    for t in configuration_tables:
        alter_column_to_timezone_aware(t, 'added')
        alter_column_to_timezone_aware(t, 'updated')
        alter_column_to_timezone_aware(t, 'destroyed')

    # AbstractResource derived tables
    resource_tables = [
        "bridge",
        "proxy",
        "smart_proxy",
        "automation_logs",
        "eotk"
    ]
    for t in resource_tables:
        alter_column_to_timezone_aware(t, 'added')
        alter_column_to_timezone_aware(t, 'updated')
        alter_column_to_timezone_aware(t, 'deprecated')
        alter_column_to_timezone_aware(t, 'destroyed')

    # Deprecation
    alter_column_to_timezone_aware("deprecation", "deprecated_at")

    # Activity
    alter_column_to_timezone_aware("activity", "added")

    # Alarm
    alter_column_to_timezone_aware("alarm", "state_changed")
    alter_column_to_timezone_aware("alarm", "last_updated")

    # Bridge terraform_updated
    alter_column_to_timezone_aware("bridge", "terraform_updated")

    # Proxy terraform_updated
    alter_column_to_timezone_aware("proxy", "terraform_updated")

    # Automation last_run, next_run
    alter_column_to_timezone_aware("automation", "last_run")
    alter_column_to_timezone_aware("automation", "next_run")

    # Onion cert_expiry
    alter_column_to_timezone_aware("onion", "cert_expiry")


def downgrade():
    pass