"""Create baseline migration

Revision ID: b93b325069b1
Revises:
Create Date: 2024-09-10 23:49:25.890085

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlalchemy_utils
from sqlalchemy import inspect

# revision identifiers, used by Alembic.
revision: str = 'b93b325069b1'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    bind = op.get_bind()
    inspector = inspect(bind)
    existing_tables = inspector.get_table_names()

    if 'wurker_control' not in existing_tables:
        op.create_table(
            'wurker_control',
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('seed', sa.Integer(), nullable=True),
            sa.Column('reload', sa.Boolean(), nullable=True),
            sa.Column('pause', sa.Boolean(), nullable=True),
            sa.PrimaryKeyConstraint('app_id')
        )

    if 'wurker_command' not in existing_tables:
        op.create_table(
            'wurker_command',
            sa.Column('id', sa.Integer(), nullable=False),
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('slug', sa.String(length=36), nullable=False),
            sa.Column('name', sa.String(length=36), nullable=False),
            sa.Column('description', sa.Text(), nullable=True),
            sa.Column('command', sa.String(length=128), nullable=True),
            sa.Column('is_module', sa.Boolean(), nullable=False),
            sa.ForeignKeyConstraint(['app_id'], ['wurker_control.app_id']),
            sa.PrimaryKeyConstraint('id'),
            sa.UniqueConstraint('app_id', 'slug')
        )

    if 'wurker_cron' not in existing_tables:
        op.create_table(
            'wurker_cron',
            sa.Column('id', sa.Integer(), nullable=False),
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('slug', sa.String(length=64), nullable=False),
            sa.Column('name', sa.String(length=64), nullable=False),
            sa.Column('cron', sa.String(length=32), nullable=True),
            sa.Column('overrun', sa.Integer(), nullable=False),
            sa.ForeignKeyConstraint(['app_id'], ['wurker_control.app_id']),
            sa.PrimaryKeyConstraint('id'),
            sa.UniqueConstraint('app_id', 'slug')
        )

    if 'wurker_status' not in existing_tables:
        op.create_table(
            'wurker_status',
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('slug', sa.String(length=36), nullable=False),
            sa.Column('name', sa.String(length=36), nullable=False),
            sa.Column('is_queen', sa.Boolean(), nullable=True),
            sa.Column('env', sa.String(length=36), nullable=False),
            sa.Column('private_ip', sa.String(length=36), nullable=True),
            sa.Column('public_ip', sa.String(length=36), nullable=True),
            sa.Column('server_domain', sa.String(length=128), nullable=True),
            sa.ForeignKeyConstraint(['app_id'], ['wurker_control.app_id']),
            sa.PrimaryKeyConstraint('app_id')
        )

    if 'wurker_bee' not in existing_tables:
        op.create_table(
            'wurker_bee',
            sa.Column('id', sa.Integer(), nullable=False),
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('slug', sa.String(length=128), nullable=False),
            sa.Column('name', sa.String(length=128), nullable=False),
            sa.Column('description', sa.Text(), nullable=True),
            sa.Column('command_id', sa.Integer(), nullable=False),
            sa.Column('args', sa.JSON(), nullable=True),
            sa.Column('quiet_restart', sa.Boolean(), nullable=True),
            sa.Column('enabled', sa.Boolean(), nullable=False),
            sa.ForeignKeyConstraint(['app_id'], ['wurker_control.app_id']),
            sa.ForeignKeyConstraint(['command_id'], ['wurker_command.id']),
            sa.PrimaryKeyConstraint('id'),
            sa.UniqueConstraint('app_id', 'slug')
        )

    if 'wurker_bee_cron' not in existing_tables:
        op.create_table(
            'wurker_bee_cron',
            sa.Column('wurker_bee_id', sa.Integer(), nullable=False),
            sa.Column('wurker_cron_id', sa.Integer(), nullable=False),
            sa.ForeignKeyConstraint(['wurker_bee_id'], ['wurker_bee.id']),
            sa.ForeignKeyConstraint(['wurker_cron_id'], ['wurker_cron.id']),
            sa.PrimaryKeyConstraint('wurker_bee_id', 'wurker_cron_id')
        )

    if 'wurker_bee_report' not in existing_tables:
        op.create_table(
            'wurker_bee_report',
            sa.Column('id', sa.Integer(), nullable=False),
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('wurker_bee_id', sa.Integer(), nullable=False),
            sa.Column('wurker_cron_id', sa.Integer(), nullable=True),
            sa.Column('hive_pid', sa.Integer(), nullable=True),
            sa.Column('hive_timestamp', sa.DateTime(timezone=True), nullable=True),
            sa.Column('hive_runcount', sa.Integer(), nullable=True),
            sa.Column('hive_duration', sa.Integer(), nullable=True),
            sa.Column('status', sa.Enum(
                'PROCESSING', 'OVERRUN', 'TIMEOUT', 'ERROR', 'COMPLETED',
                name='wurkerbeestate', create_type=False
            ), nullable=False),
            sa.Column('status_info', sa.JSON(), nullable=True),
            sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=True),
            sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=True),
            sa.ForeignKeyConstraint(['app_id'], ['wurker_control.app_id'], ondelete='CASCADE'),
            sa.ForeignKeyConstraint(['wurker_bee_id'], ['wurker_bee.id'], ondelete='CASCADE'),
            sa.ForeignKeyConstraint(['wurker_cron_id'], ['wurker_cron.id'], ondelete='CASCADE'),
            sa.PrimaryKeyConstraint('id')
        )

    if 'wurker_queue' not in existing_tables:
        op.create_table(
            'wurker_queue',
            sa.Column('wurker_bee_id', sa.Integer(), nullable=False),
            sa.Column('app_id', sqlalchemy_utils.types.uuid.UUIDType(), nullable=False),
            sa.Column('args', sa.JSON(), nullable=True),
            sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=True),
            sa.ForeignKeyConstraint(['app_id'], ['wurker_control.app_id']),
            sa.ForeignKeyConstraint(['wurker_bee_id'], ['wurker_bee.id']),
            sa.PrimaryKeyConstraint('wurker_bee_id')
        )


def downgrade() -> None:
    bind = op.get_bind()
    inspector = inspect(bind)
    existing_tables = inspector.get_table_names()

    if 'wurker_queue' in existing_tables:
        op.drop_table('wurker_queue')
    if 'wurker_bee_report' in existing_tables:
        op.drop_table('wurker_bee_report')
    if 'wurker_bee_cron' in existing_tables:
        op.drop_table('wurker_bee_cron')
    if 'wurker_bee' in existing_tables:
        op.drop_table('wurker_bee')
    if 'wurker_status' in existing_tables:
        op.drop_table('wurker_status')
    if 'wurker_cron' in existing_tables:
        op.drop_table('wurker_cron')
    if 'wurker_command' in existing_tables:
        op.drop_table('wurker_command')
    if 'wurker_control' in existing_tables:
        op.drop_table('wurker_control')

    if bind.dialect.name == 'postgresql':
        op.execute("DROP TYPE IF EXISTS wurkerbeestate")
