"""Automations models

Revision ID: aeea5ee6f070
Revises: 7a653837d9ba
Create Date: 2024-04-03 11:24:09.599513

"""

from typing import List

import sqlalchemy as sa
from alembic import op

import prefect
from prefect.server.events.actions import ActionTypes
from prefect.server.events.schemas.automations import Firing, TriggerTypes
from prefect.server.events.schemas.events import ReceivedEvent

# revision identifiers, used by Alembic.
revision = "aeea5ee6f070"
down_revision = "7a653837d9ba"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "automation",
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=False),
        sa.Column("enabled", sa.Boolean(), server_default="1", nullable=False),
        sa.Column(
            "trigger",
            prefect.server.utilities.database.Pydantic(TriggerTypes),
            nullable=False,
        ),
        sa.Column(
            "actions",
            prefect.server.utilities.database.Pydantic(List[ActionTypes]),
            nullable=False,
        ),
        sa.Column(
            "actions_on_trigger",
            prefect.server.utilities.database.Pydantic(List[ActionTypes]),
            server_default="[]",
            nullable=False,
        ),
        sa.Column(
            "actions_on_resolve",
            prefect.server.utilities.database.Pydantic(List[ActionTypes]),
            server_default="[]",
            nullable=False,
        ),
        sa.Column(
            "id",
            prefect.server.utilities.database.UUID(),
            server_default=sa.text("(GEN_RANDOM_UUID())"),
            nullable=False,
        ),
        sa.Column(
            "created",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.Column(
            "updated",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_automation")),
    )
    op.create_index(
        op.f("ix_automation__updated"), "automation", ["updated"], unique=False
    )
    op.create_table(
        "automation_bucket",
        sa.Column(
            "automation_id", prefect.server.utilities.database.UUID(), nullable=False
        ),
        sa.Column(
            "trigger_id", prefect.server.utilities.database.UUID(), nullable=True
        ),
        sa.Column(
            "bucketing_key",
            prefect.server.utilities.database.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column(
            "last_event",
            prefect.server.utilities.database.Pydantic(ReceivedEvent),
            nullable=True,
        ),
        sa.Column(
            "start",
            prefect.server.utilities.database.Timestamp(timezone=True),
            nullable=False,
        ),
        sa.Column(
            "end",
            prefect.server.utilities.database.Timestamp(timezone=True),
            nullable=False,
        ),
        sa.Column("count", sa.Integer(), nullable=False),
        sa.Column("last_operation", sa.String(), nullable=True),
        sa.Column(
            "triggered_at",
            prefect.server.utilities.database.Timestamp(timezone=True),
            nullable=False,
        ),
        sa.Column(
            "id",
            prefect.server.utilities.database.UUID(),
            server_default=sa.text("(GEN_RANDOM_UUID())"),
            nullable=False,
        ),
        sa.Column(
            "created",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.Column(
            "updated",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["automation_id"],
            ["automation.id"],
            name=op.f("fk_automation_bucket__automation_id__automation"),
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_automation_bucket")),
    )
    op.create_index(
        "ix_automation_bucket__automation_id__end",
        "automation_bucket",
        ["automation_id", "end"],
        unique=False,
    )
    op.create_index(
        op.f("ix_automation_bucket__updated"),
        "automation_bucket",
        ["updated"],
        unique=False,
    )
    op.create_index(
        "uq_automation_bucket__automation_id__bucketing_key",
        "automation_bucket",
        ["automation_id", "bucketing_key"],
        unique=True,
    )
    op.create_table(
        "automation_related_resource",
        sa.Column(
            "automation_id", prefect.server.utilities.database.UUID(), nullable=False
        ),
        sa.Column("resource_id", sa.String(), nullable=True),
        sa.Column(
            "automation_owned_by_resource",
            sa.Boolean(),
            server_default="0",
            nullable=False,
        ),
        sa.Column(
            "id",
            prefect.server.utilities.database.UUID(),
            server_default=sa.text("(GEN_RANDOM_UUID())"),
            nullable=False,
        ),
        sa.Column(
            "created",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.Column(
            "updated",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["automation_id"],
            ["automation.id"],
            name=op.f("fk_automation_related_resource__automation_id__automation"),
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_automation_related_resource")),
    )
    op.create_index(
        op.f("ix_automation_related_resource__resource_id"),
        "automation_related_resource",
        ["resource_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_automation_related_resource__updated"),
        "automation_related_resource",
        ["updated"],
        unique=False,
    )
    op.create_index(
        "uq_automation_related_resource__automation_id__resource_id",
        "automation_related_resource",
        ["automation_id", "resource_id"],
        unique=True,
    )
    op.create_table(
        "composite_trigger_child_firing",
        sa.Column(
            "automation_id", prefect.server.utilities.database.UUID(), nullable=False
        ),
        sa.Column(
            "parent_trigger_id",
            prefect.server.utilities.database.UUID(),
            nullable=False,
        ),
        sa.Column(
            "child_trigger_id", prefect.server.utilities.database.UUID(), nullable=False
        ),
        sa.Column(
            "child_firing_id", prefect.server.utilities.database.UUID(), nullable=False
        ),
        sa.Column(
            "child_fired_at",
            prefect.server.utilities.database.Timestamp(timezone=True),
            nullable=True,
        ),
        sa.Column(
            "child_firing",
            prefect.server.utilities.database.Pydantic(Firing),
            nullable=False,
        ),
        sa.Column(
            "id",
            prefect.server.utilities.database.UUID(),
            server_default=sa.text("(GEN_RANDOM_UUID())"),
            nullable=False,
        ),
        sa.Column(
            "created",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.Column(
            "updated",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("CURRENT_TIMESTAMP"),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["automation_id"],
            ["automation.id"],
            name=op.f("fk_composite_trigger_child_firing__automation_id__automation"),
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_composite_trigger_child_firing")),
    )
    op.create_index(
        op.f("ix_composite_trigger_child_firing__updated"),
        "composite_trigger_child_firing",
        ["updated"],
        unique=False,
    )
    op.create_index(
        "uq_composite_trigger_child_firing__a_id__pt_id__ct__id",
        "composite_trigger_child_firing",
        ["automation_id", "parent_trigger_id", "child_trigger_id"],
        unique=True,
    )


def downgrade():
    op.drop_index(
        "uq_composite_trigger_child_firing__a_id__pt_id__ct__id",
        table_name="composite_trigger_child_firing",
    )
    op.drop_index(
        op.f("ix_composite_trigger_child_firing__updated"),
        table_name="composite_trigger_child_firing",
    )
    op.drop_table("composite_trigger_child_firing")
    op.drop_index(
        "uq_automation_related_resource__automation_id__resource_id",
        table_name="automation_related_resource",
    )
    op.drop_index(
        op.f("ix_automation_related_resource__updated"),
        table_name="automation_related_resource",
    )
    op.drop_index(
        op.f("ix_automation_related_resource__resource_id"),
        table_name="automation_related_resource",
    )
    op.drop_table("automation_related_resource")
    op.drop_index(
        "uq_automation_bucket__automation_id__bucketing_key",
        table_name="automation_bucket",
    )
    op.drop_index(op.f("ix_automation_bucket__updated"), table_name="automation_bucket")
    op.drop_index(
        "ix_automation_bucket__automation_id__end", table_name="automation_bucket"
    )
    op.drop_table("automation_bucket")
    op.drop_index(op.f("ix_automation__updated"), table_name="automation")
    op.drop_table("automation")
    # ### end Alembic commands ###
