# Copyright (c) 2016 Mirantis Inc.
# 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.

"""Refactor task results

Revision ID: e654a0648db0
Revises: 3177d36ea270
Create Date: 2016-04-01 14:36:56.373349

"""

import datetime as dt
import json
import uuid

from alembic import op
import sqlalchemy as sa

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


# revision identifiers, used by Alembic.
revision = "e654a0648db0"
down_revision = "32fada9b2fde"
branch_labels = None
depends_on = None


taskhelper = 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(
        "status",
        sa.Enum(
            "aborted",
            "aborting",
            "cleaning up",
            "failed",
            "finished",
            "init",
            "paused",
            "running",
            "setting up",
            "soft_aborting",
            "verifying",
            name="enum_tasks_status",
        ),
        nullable=False,
    ),
    sa.Column("verification_log", sa.Text(), nullable=True),
    sa.Column("tag", sa.String(length=64), nullable=True),
    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(
        "validation_result",
        sa_types.MutableJSONEncodedDict(),
        default={},
        nullable=False,
    ),
)

task_result_helper = sa.Table(
    "task_results",
    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("key", sa_types.MutableJSONEncodedDict(), nullable=False),
    sa.Column("data", sa_types.MutableJSONEncodedDict(), nullable=False),
    sa.Column("task_uuid", sa.String(length=36), nullable=True),
)

taghelper = sa.Table(
    "tags",
    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("tag", sa.String(length=255), nullable=False),
    sa.Column(
        "type",
        sa.Enum("task", "subtask", name="enum_tag_types"),
        nullable=False,
    ),
)


def upgrade() -> None:
    conn = op.get_bind()

    subtask_table = op.create_table(
        "subtasks",
        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",
            sa.Enum(
                "finished", "running", "crashed", name="enum_subtasks_status"
            ),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["task_uuid"],
            ["tasks.uuid"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_index("subtask_uuid", "subtasks", ["uuid"], unique=True)
    op.create_index("subtask_status", "subtasks", ["status"], unique=False)

    workload_table = op.create_table(
        "workloads",
        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("subtask_uuid", sa.String(length=36), nullable=False),
        sa.Column("name", sa.String(length=64), nullable=False),
        sa.Column("description", sa.Text(), default=""),
        sa.Column("position", sa.Integer(), default=0, nullable=False),
        sa.Column("runner_type", sa.String(length=64), nullable=False),
        sa.Column(
            "runner",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column(
            "args",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column(
            "context",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column(
            "hooks",
            sa_types.MutableJSONEncodedList(),
            default=[],
            nullable=False,
        ),
        sa.Column(
            "sla",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column(
            "sla_results",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column(
            "context_execution",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column("load_duration", sa.Float(), default=0),
        sa.Column("full_duration", sa.Float(), default=0),
        sa.Column("min_duration", sa.Float(), default=0),
        sa.Column("max_duration", sa.Float(), default=0),
        sa.Column("total_iteration_count", sa.Integer(), default=0),
        sa.Column("failed_iteration_count", sa.Integer(), default=0),
        sa.Column("pass_sla", sa.Boolean()),
        sa.Column(
            "statistics",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.Column("start_time", sa.DateTime()),
        sa.Column("_profiling_data", sa.Text(), default=""),
        sa.ForeignKeyConstraint(
            ["task_uuid"],
            ["tasks.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["subtask_uuid"],
            ["subtasks.uuid"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_index("workload_uuid", "workloads", ["uuid"], unique=True)

    workloaddata_table = op.create_table(
        "workloaddata",
        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("workload_uuid", sa.String(length=36), nullable=False),
        sa.Column("chunk_order", sa.Integer(), nullable=False),
        sa.Column("iteration_count", sa.Integer(), nullable=False),
        sa.Column("failed_iteration_count", sa.Integer(), nullable=False),
        sa.Column("chunk_size", sa.Integer(), nullable=False),
        sa.Column("compressed_chunk_size", sa.Integer(), nullable=False),
        sa.Column("started_at", sa.DateTime(), nullable=False),
        sa.Column("finished_at", sa.DateTime(), nullable=False),
        sa.Column(
            "chunk_data",
            sa_types.MutableJSONEncodedDict(),
            default={},
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["task_uuid"],
            ["tasks.uuid"],
        ),
        sa.ForeignKeyConstraint(
            ["workload_uuid"],
            ["workloads.uuid"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_index(
        "workload_data_uuid", "workloaddata", ["uuid"], unique=True
    )

    tag_table = op.create_table(
        "tags",
        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("tag", sa.String(length=255), nullable=False),
        sa.Column(
            "type",
            sa.Enum("task", "subtask", name="enum_tag_types"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_index("d_type_tag", "tags", ["uuid", "type", "tag"], unique=True)

    with op.batch_alter_table("tasks") as batch_op:
        batch_op.add_column(
            sa.Column("title", sa.String(length=64), default="")
        )

        batch_op.add_column(sa.Column("description", sa.Text(), default=""))

        batch_op.add_column(sa.Column("input_task", sa.Text(), default=""))

        batch_op.add_column(sa.Column("validation_duration", sa.Float()))

        batch_op.add_column(sa.Column("task_duration", sa.Float()))

        batch_op.add_column(sa.Column("pass_sla", sa.Boolean()))

        batch_op.add_column(
            sa.Column(
                "validation_result",
                sa_types.MutableJSONEncodedDict(),
                default={},
            )
        )

    for task in conn.execute(taskhelper.select()):
        if task.tag:
            conn.execute(
                tag_table.insert(),
                [
                    {
                        "uuid": task.uuid,
                        "type": "task",
                        "tag": task.tag,
                        "created_at": task.created_at,
                        "updated_at": task.updated_at,
                    }
                ],
            )

        task_results = conn.execute(
            task_result_helper.select().where(
                task_result_helper.c.task_uuid == task.uuid
            )
        )

        pass_sla = True
        task_duration = 0

        for task_result in task_results:
            raw_data = task_result.data.get("raw", [])
            iter_count = len(raw_data)

            failed_iter_count = 0
            max_duration = 0
            min_duration = -1

            for d in raw_data:
                if d.get("error"):
                    failed_iter_count += 1

                duration = d.get("duration", 0)

                if duration > max_duration:
                    max_duration = duration

                if min_duration < 0 or min_duration > duration:
                    min_duration = duration

            sla = task_result.data.get("sla", [])
            success = all(s.get("success") for s in sla)

            if not success:
                pass_sla = False

            task_duration += task_result.data.get("full_duration", 0)

            delta = dt.timedelta(
                seconds=task_result.data.get("full_duration", 0)
            )
            start = task_result.created_at - delta

            subtask_uuid = str(uuid.uuid4())

            conn.execute(
                subtask_table.insert(),
                [
                    {
                        "uuid": subtask_uuid,
                        "task_uuid": task.uuid,
                        "created_at": task_result.created_at,
                        "updated_at": task_result.updated_at,
                        # NOTE(ikhudoshyn) We don't have info on subtask status
                        "status": "finished",
                        "duration": task_result.data.get("full_duration", 0),
                        "pass_sla": success,
                    }
                ],
            )

            workload_uuid = str(uuid.uuid4())

            conn.execute(
                workload_table.insert(),
                [
                    {
                        "created_at": task_result.created_at,
                        "updated_at": task_result.updated_at,
                        "uuid": workload_uuid,
                        "task_uuid": task.uuid,
                        "subtask_uuid": subtask_uuid,
                        "name": task_result.key["name"],
                        "position": task_result.key["pos"],
                        "runner_type": task_result.key["kw"]["runner"]["type"],
                        "runner": task_result.key["kw"]["runner"],
                        "context": task_result.key["kw"].get("context", {}),
                        "sla": task_result.key["kw"].get("sla", {}),
                        "args": task_result.key["kw"].get("args", {}),
                        "sla_results": {"sla": sla},
                        "context_execution": {},
                        "load_duration": task_result.data.get(
                            "load_duration", 0
                        ),
                        "full_duration": task_result.data.get(
                            "full_duration", 0
                        ),
                        "min_duration": min_duration,
                        "max_duration": max_duration,
                        "total_iteration_count": iter_count,
                        "failed_iteration_count": failed_iter_count,
                        "pass_sla": success,
                        "statistics": {},
                        "start_time": start,
                    }
                ],
            )

            conn.execute(
                workloaddata_table.insert(),
                [
                    {
                        "uuid": str(uuid.uuid4()),
                        "task_uuid": task.uuid,
                        "workload_uuid": workload_uuid,
                        "chunk_order": 0,
                        "iteration_count": iter_count,
                        "failed_iteration_count": failed_iter_count,
                        "chunk_data": {"raw": raw_data},
                        # TODO(ikhudoshyn)
                        "chunk_size": 0,
                        "compressed_chunk_size": 0,
                        "started_at": start,
                        "finished_at": task_result.created_at,
                    }
                ],
            )

        task_verification_log = {}
        if task.verification_log:
            task_verification_log = json.loads(task.verification_log)

        conn.execute(
            taskhelper.update().where(taskhelper.c.uuid == task.uuid),
            {
                "pass_sla": pass_sla,
                "task_duration": task_duration,
                "validation_duration": 0,
                "validation_result": task_verification_log,
            },
        )

    # TODO(ikhudoshyn) update workload's statistics

    with op.batch_alter_table("tasks") as batch_op:
        batch_op.drop_column("tag")
        batch_op.drop_column("verification_log")
        batch_op.alter_column(
            "validation_result",
            existing_type=sa_types.MutableJSONEncodedDict(),
            nullable=False,
        )

    op.drop_table("task_results")


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