# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""Change task status enum

Revision ID: 4ef544102ba7
Revises: 3177d36ea270
Create Date: 2016-04-22 21:28:50.745316

"""

from alembic import op
import sqlalchemy as sa

from rally import consts
from rally import exceptions
from rally.common.db import sa_types


# revision identifiers, used by Alembic.
revision = "4ef544102ba7"
down_revision = "f33f4610dcda"
branch_labels = None
depends_on = None


OLD_STATUS = [
    "aborted",
    "aborting",
    "cleaning up",
    "failed",
    "finished",
    "init",
    "paused",
    "running",
    "setting up",
    "soft_aborting",
    "verifying",
]
OLD_ENUM = sa.Enum(*OLD_STATUS, name="enum_tasks_status")

WITHOUT_CHANGES = (
    "init",
    "running",
    "aborted",
    "aborting",
    "soft_aborting",
    "paused",
    "finished",
)

OLD_TO_NEW = [
    (
        "verifying",
        "validating",
    ),
    (
        "failed",
        "crashed",
    ),
]


task = sa.Table(
    "tasks",
    sa.MetaData(),
    sa.Column("created_at", sa.DateTime(), nullable=True),
    sa.Column("updated_at", sa.DateTime(), nullable=True),
    sa.Column("id", sa.Integer(), nullable=False),
    sa.Column("uuid", sa.String(length=36), nullable=False),
    sa.Column("deployment_uuid", sa.String(length=36), nullable=False),
    sa.Column("title", sa.String(length=64), default=""),
    sa.Column("description", sa.Text(), default=""),
    sa.Column("input_task", sa.Text(), default=""),
    sa.Column("validation_duration", sa.Float()),
    sa.Column("task_duration", sa.Float()),
    sa.Column("pass_sla", sa.Boolean()),
    sa.Column("status", OLD_ENUM, nullable=False),
    sa.Column("new_status", sa.String(36), default=consts.TaskStatus.INIT),
    sa.Column(
        "validation_result",
        sa_types.MutableJSONEncodedDict(),
        default={},
        nullable=False,
    ),
)

subtask = sa.Table(
    "subtasks",
    sa.MetaData(),
    sa.Column("created_at", sa.DateTime()),
    sa.Column("updated_at", sa.DateTime()),
    sa.Column("id", sa.Integer(), nullable=False, autoincrement=True),
    sa.Column("uuid", sa.String(length=36), nullable=False),
    sa.Column("task_uuid", sa.String(length=36), nullable=False),
    sa.Column("title", sa.String(length=64), default=""),
    sa.Column("description", sa.Text(), default=""),
    sa.Column(
        "context",
        sa_types.MutableJSONEncodedDict(),
        default={},
        nullable=False,
    ),
    sa.Column(
        "sla", sa_types.MutableJSONEncodedDict(), default={}, nullable=False
    ),
    sa.Column("duration", sa.Float()),
    sa.Column("run_in_parallel", sa.Boolean(), default=False, nullable=False),
    sa.Column("pass_sla", sa.Boolean()),
    sa.Column("status", OLD_ENUM, nullable=False),
    sa.Column(
        "new_status", sa.String(36), default=consts.SubtaskStatus.RUNNING
    ),
    sa.ForeignKeyConstraint(
        ["task_uuid"],
        ["tasks.uuid"],
    ),
    sa.PrimaryKeyConstraint("id"),
)


def upgrade() -> None:
    # Workaround for Alemic bug #89
    # https://bitbucket.org/zzzeek/alembic/issue/89

    with op.batch_alter_table("tasks") as batch_op:
        batch_op.add_column(
            sa.Column(
                "new_status", sa.String(36), default=consts.TaskStatus.INIT
            )
        )
    with op.batch_alter_table("subtasks") as batch_op:
        batch_op.add_column(
            sa.Column(
                "new_status",
                sa.String(36),
                default=consts.SubtaskStatus.RUNNING,
            )
        )

    op.execute(
        task.update()
        .where(task.c.status.in_(WITHOUT_CHANGES))
        .values({"new_status": task.c.status})
    )

    for old, new in OLD_TO_NEW:
        op.execute(
            task.update()
            .where(task.c.status == op.inline_literal(old))
            .values({"new_status": new})
        )

    # NOTE(rvasilets): Assume that set_failed was used only in causes of
    # validation failed
    op.execute(
        task.update()
        .where(
            (task.c.status == op.inline_literal("failed"))
            & (task.c.validation_result == {})
        )
        .values({"new_status": "crashed", "validation_result": {}})
    )
    op.execute(
        task.update()
        .where(
            (task.c.status == op.inline_literal("failed"))
            & (task.c.validation_result != {})
        )
        .values(
            {
                "new_status": "validation_failed",
                "validation_result": task.c.validation_result,
            }
        )
    )

    op.drop_index("task_status", "tasks")
    op.drop_index("subtask_status", "subtasks")

    # NOTE(boris-42): Statuses "setting up", "cleaning up" were not used
    with op.batch_alter_table("tasks") as batch_op:
        batch_op.drop_column("status")
        batch_op.alter_column(
            "new_status", new_column_name="status", existing_type=sa.String(36)
        )
    with op.batch_alter_table("subtasks") as batch_op:
        batch_op.drop_column("status")
        batch_op.alter_column(
            "new_status", new_column_name="status", existing_type=sa.String(36)
        )

    op.create_index("task_status", "tasks", ["status"])
    op.create_index("subtask_status", "subtasks", ["status"])


def downgrade() -> None:
    raise exceptions.DowngradeNotSupported()
