"""Migrate services

Revision ID: 651954cbdec6
Revises: 98aea37d208d
Create Date: 2021-07-13 09:53:23.243644

"""

import json

import numpy
import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm.session import Session
from sqlalchemy.sql import table, column

from qcfractal.db_socket.column_types import MsgpackExt, PlainMsgpackExt

# revision identifiers, used by Alembic.
revision = "651954cbdec6"
down_revision = "98aea37d208d"
branch_labels = None
depends_on = None


def convert_numpy(d):
    if isinstance(d, dict):
        for k, v in d.items():
            d[k] = convert_numpy(v)
        return d
    elif isinstance(d, numpy.ndarray):
        return d.tolist()
    elif isinstance(d, list):
        return [convert_numpy(x) for x in d]
    else:
        return d


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###

    # Create the auxiliary table for tasks associated with services
    op.create_table(
        "service_queue_tasks",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("service_id", sa.Integer(), nullable=False),
        sa.Column("procedure_id", sa.Integer(), nullable=False),
        sa.Column("extras", sa.JSON(), nullable=True),
        sa.ForeignKeyConstraint(["procedure_id"], ["base_result.id"], ondelete="cascade"),
        sa.ForeignKeyConstraint(["service_id"], ["service_queue.id"], ondelete="cascade"),
        sa.PrimaryKeyConstraint("id"),
    )

    # Handle the service state
    # Add the column and then migrate what was in extra
    op.add_column("service_queue", sa.Column("service_state", PlainMsgpackExt(), nullable=True))

    # Temporary ORM
    service_tasks_table = table(
        "service_queue_tasks",
        column("service_id", sa.Integer),
        column("procedure_id", sa.Integer),
        column("extras", sa.JSON),
    )

    service_table = table(
        "service_queue",
        column("extra", MsgpackExt),
        column("service_state", PlainMsgpackExt),
        column("id", sa.Integer),
        column("tag", sa.String),
        column("priority", sa.Integer),
        column("procedure_id", sa.Integer),
    )

    kvstore_table = table("kv_store", column("id", sa.Integer), column("value", sa.JSON))

    baseresult_table = table(
        "base_result", column("id", sa.Integer), column("stdout", sa.Integer), column("error", sa.Integer)
    )

    bind = op.get_bind()
    session = Session(bind=bind)

    services = session.query(service_table).all()

    for service in services:
        if "torsiondrive_state" in service.extra:
            # We have a torsiondrive
            service_state = {}
            service_state["torsiondrive_state"] = convert_numpy(service.extra["torsiondrive_state"])

            service_state["molecule_template"] = service.extra["molecule_template"]
            service_state["dihedral_template"] = service.extra["dihedral_template"]

            old_opt_template = json.loads(service.extra["optimization_template"])
            old_opt_template["meta"]["priority"] = service.priority

            # We don't use "meta" anymore. Just store the value
            service_state["optimization_template"] = json.dumps(old_opt_template["meta"])

            # Pull out the maps of subtasks, and convert them to the new association table
            for td_api_key, task_keys in service.extra["task_map"].items():
                for task_key in task_keys:
                    task_id = service.extra["task_manager"]["required_tasks"][task_key]
                    position = task_key.split("-")[-1]

                    insert_values = {
                        "service_id": int(service.id),
                        "procedure_id": int(task_id),
                        "extras": {"td_api_key": td_api_key, "position": position},
                    }
                    session.execute(service_tasks_table.insert().values(insert_values))

        else:
            # We have a grid optimization
            service_state = {}
            service_state["iteration"] = service.extra["iteration"]
            service_state["dimensions"] = service.extra["dimensions"]
            service_state["complete"] = service.extra["complete"]
            service_state["constraint_template"] = service.extra["constraint_template"]

            old_opt_template = json.loads(service.extra["optimization_template"])
            old_opt_template["meta"]["priority"] = service.priority

            # We don't use "meta" anymore. Just store the value
            service_state["optimization_template"] = json.dumps(old_opt_template["meta"])

            # Pull out the maps of subtasks, and convert them to the new association table
            for key, task_id in service.extra["task_manager"]["required_tasks"].items():
                insert_values = {"service_id": int(service.id), "procedure_id": int(task_id), "extras": {"key": key}}

                session.execute(service_tasks_table.insert().values(insert_values))

        # Set these to be the same priority as the tasks. I don't believe the priority or tag
        # of the service themselves was ever modifiable
        update_values = {
            "service_state": service_state,
            "tag": service.extra["task_tag"],
            "priority": service.extra["task_priority"],
        }

        session.execute(service_table.update().values(update_values).where(service_table.c.id == service.id))

        # Handle outputs and errors
        stdout = service.extra.get("stdout")
        if stdout:
            stdout_id = session.execute(kvstore_table.insert().values({"value": stdout}).returning(sa.column("id")))
            stdout_id = list(stdout_id)
            assert len(stdout_id) == 1 and len(stdout_id[0]) == 1
            stmt = (
                baseresult_table.update()
                .values({"stdout": stdout_id[0][0]})
                .where(baseresult_table.c.id == service.procedure_id)
            )
            session.execute(stmt)

        error = service.extra.get("error")
        if error:
            error_id = session.execute(kvstore_table.insert().values({"value": error}).returning(sa.column("id")))
            error_id = list(error_id)
            assert len(error_id) == 1 and len(error_id[0]) == 1
            stmt = (
                baseresult_table.update()
                .values({"error": error_id[0][0]})
                .where(baseresult_table.c.id == service.procedure_id)
            )
            session.execute(stmt)

        session.flush()

    # Now drop the unneeded stuff
    op.drop_index("ix_grid_optmization_program", table_name="grid_optimization_procedure")
    op.drop_column("grid_optimization_procedure", "program")
    op.drop_index("ix_service_queue_hash_index", table_name="service_queue")
    op.drop_column("service_queue", "modified_on")
    op.drop_column("service_queue", "extra")
    op.drop_column("service_queue", "hash_index")
    op.drop_index("ix_torsion_drive_program", table_name="torsiondrive_procedure")
    op.drop_column("torsiondrive_procedure", "program")
    # ### end Alembic commands ###


def downgrade():
    raise RuntimeError("Cannot downgrade")

    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "torsiondrive_procedure", sa.Column("program", sa.VARCHAR(length=100), autoincrement=False, nullable=False)
    )
    op.create_index("ix_torsion_drive_program", "torsiondrive_procedure", ["program"], unique=False)
    op.add_column("service_queue", sa.Column("hash_index", sa.VARCHAR(), autoincrement=False, nullable=False))
    op.add_column("service_queue", sa.Column("extra", postgresql.BYTEA(), autoincrement=False, nullable=True))
    op.add_column("service_queue", sa.Column("modified_on", postgresql.TIMESTAMP(), autoincrement=False, nullable=True))
    op.create_index("ix_service_queue_hash_index", "service_queue", ["hash_index"], unique=False)
    op.drop_column("service_queue", "service_state")
    op.add_column(
        "grid_optimization_procedure", sa.Column("program", sa.VARCHAR(length=100), autoincrement=False, nullable=False)
    )
    op.create_index("ix_grid_optmization_program", "grid_optimization_procedure", ["program"], unique=False)
    op.drop_table("service_queue_tasks")
    # ### end Alembic commands ###
