"""refactor user expire

Revision ID: a83e4dd22672
Revises: 9001356f0063
Create Date: 2024-06-27 01:04:34.974520

"""

import sqlalchemy as sa
from alembic import op
from sqlalchemy import table, column
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "a83e4dd22672"
down_revision = "9001356f0063"
branch_labels = None
depends_on = None


def upgrade() -> None:
    bind = op.get_bind()
    dialect = bind.dialect.name

    if dialect == "postgresql":
        op.execute(
            """
            DO $$ 
            BEGIN 
                IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'userexpirestrategy') THEN
                    CREATE TYPE userexpirestrategy AS ENUM ('NEVER', 'FIXED_DATE', 'START_ON_FIRST_USE');
                END IF;
            END $$;
            """
        )
    op.add_column(
        "users",
        sa.Column(
            "activated",
            sa.Boolean(),
            nullable=False,
            server_default=sa.sql.true(),
        ),
    )
    op.add_column(
        "users",
        sa.Column(
            "expire_strategy",
            sa.Enum(
                "NEVER",
                "FIXED_DATE",
                "START_ON_FIRST_USE",
                name="userexpirestrategy",
            ),
            server_default="NEVER",
            nullable=False,
        ),
    )

    if dialect == "sqlite":
        for old, new, type_ in [
            ("on_hold_timeout", "activation_deadline", sa.DateTime),
            ("expire", "expire_date", sa.DateTime),
            ("on_hold_expire_duration", "usage_duration", sa.Integer),
        ]:
            op.add_column("users", sa.Column(new, type_))

            # Step 2: Copy data from the old column to the new column
            op.execute(f"UPDATE users SET {new} = {old}")

            with op.batch_alter_table("users") as batch_op:
                batch_op.drop_column(old)
    else:
        op.alter_column(
            "users",
            "on_hold_timeout",
            new_column_name="activation_deadline",
            existing_type=sa.DateTime,
        )
        op.alter_column(
            "users",
            "expire",
            new_column_name="expire_date",
            existing_type=sa.DateTime,
        )
        op.alter_column(
            "users",
            "on_hold_expire_duration",
            new_column_name="usage_duration",
            existing_type=sa.BigInteger,
        )
    # ### end Alembic commands ###
    users_table = table(
        "users",
        column("id", sa.Integer),
        column("status", sa.String),
        column("expire_date", sa.DateTime),
        column("expire_strategy", sa.Enum(name="userexpirestrategy")),
    )

    connection = op.get_bind()
    results = connection.execute(
        sa.select(
            users_table.c.id, users_table.c.status, users_table.c.expire_date
        )
    )

    for user in results:
        status = user[1]
        expire_date = user[2]

        if status == "on_hold":
            expire_strategy = "START_ON_FIRST_USE"
        elif expire_date:
            expire_strategy = "FIXED_DATE"
        else:
            expire_strategy = "NEVER"
        connection.execute(
            users_table.update()
            .where(users_table.c.id == user[0])
            .values(expire_strategy=expire_strategy)
        )

    op.drop_column("users", "status")


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "users",
        sa.Column(
            "status",
            postgresql.ENUM(
                "active", "limited", "expired", "on_hold", name="userstatus"
            ),
            autoincrement=False,
            nullable=False,
            server_default="active",
        ),
    )
    op.alter_column(
        "users",
        "activation_deadline",
        new_column_name="on_hold_timeout",
        existing_type=sa.DateTime,
    )
    op.alter_column(
        "users",
        "usage_duration",
        new_column_name="on_hold_expire_duration",
        existing_type=sa.BigInteger,
    )
    op.alter_column(
        "users",
        "expire_date",
        new_column_name="expire",
        existing_type=sa.DateTime,
    )
    op.drop_column("users", "expire_strategy")
    op.drop_column("users", "activated")
    # ### end Alembic commands ###