"""Migrate task queue orm spec

Revision ID: 11a500398106
Revises: c0fba229aaa8
Create Date: 2023-01-17 20:24:28.657641

"""

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

from qcfractal.db_socket.column_types import MsgpackExt

# revision identifiers, used by Alembic.
revision = "11a500398106"
down_revision = "c0fba229aaa8"
branch_labels = None
depends_on = None


def convert_numpy(obj):
    if isinstance(obj, dict):
        return {k: convert_numpy(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple, set)):
        return [convert_numpy(v) for v in obj]
    elif isinstance(obj, np.ndarray):
        if obj.shape:
            return obj.ravel().tolist()
        else:
            return obj.tolist()
    else:
        return obj


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("task_queue", sa.Column("function", sa.String(), nullable=True))
    op.add_column("task_queue", sa.Column("function_kwargs", sa.JSON(), nullable=True))

    task_table = table(
        "task_queue",
        column("id", sa.Integer),
        column("spec", MsgpackExt),
        column("function", sa.String),
        column("function_kwargs", sa.JSON),
    )

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

    tasks = session.query(task_table).where(task_table.c.spec.is_not(None)).yield_per(200)

    for task in tasks:
        if task.spec["function"] == "qcengine.compute":
            # Singlepoint
            function_kwargs = {
                "program": task.spec["args"][1],
                "input_data": convert_numpy(task.spec["args"][0]),
            }

        elif task.spec["function"] == "qcengine.compute_procedure":
            function_kwargs = {
                "procedure": task.spec["args"][1],
                "input_data": convert_numpy(task.spec["args"][0]),
            }
        else:
            raise RuntimeError(f"Unknown task function: {task.spec['function']}")

        # Pop some extras if specified
        if "input_data" in function_kwargs and "extras" in function_kwargs["input_data"]:
            function_kwargs["input_data"]["extras"].pop("_qcfractal_tags", None)

        session.execute(
            sa.update(task_table)
            .values(function=task.spec["function"], function_kwargs=function_kwargs)
            .where(task_table.c.id == task.id)
        )

    op.drop_column("task_queue", "spec")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    raise RuntimeError("Cannot downgrade")
    # ### end Alembic commands ###
