"""Add collection tables (dataset, reaction_dataset)

Revision ID: 129ff3ce9247
Revises: e32b61e2516f
Create Date: 2019-08-19 16:35:19.100113

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.orm.session import Session

from qcfractal.storage_sockets.models import CollectionORM, DatasetORM, ReactionDatasetORM


# revision identifiers, used by Alembic.
revision = "129ff3ce9247"
down_revision = "e32b61e2516f"
branch_labels = None
depends_on = None


def upgrade():
    print("Start schema migration...")
    ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "dataset",
        sa.Column("default_benchmark", sa.String(), nullable=True),
        sa.Column("default_keywords", sa.JSON(), nullable=True),
        sa.Column("default_driver", sa.String(), nullable=True),
        sa.Column("default_units", sa.String(), nullable=True),
        sa.Column("alias_keywords", sa.JSON(), nullable=True),
        sa.Column("default_program", sa.String(), nullable=True),
        sa.Column("contributed_values", sa.JSON(), nullable=True),
        sa.Column("history_keys", sa.JSON(), nullable=True),
        sa.Column("history", sa.JSON(), nullable=True),
        sa.Column("id", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(["id"], ["collection.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "reaction_dataset",
        sa.Column("default_benchmark", sa.String(), nullable=True),
        sa.Column("default_keywords", sa.JSON(), nullable=True),
        sa.Column("default_driver", sa.String(), nullable=True),
        sa.Column("default_units", sa.String(), nullable=True),
        sa.Column("alias_keywords", sa.JSON(), nullable=True),
        sa.Column("default_program", sa.String(), nullable=True),
        sa.Column("contributed_values", sa.JSON(), nullable=True),
        sa.Column("history_keys", sa.JSON(), nullable=True),
        sa.Column("history", sa.JSON(), nullable=True),
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("ds_type", sa.String(), nullable=True),
        sa.ForeignKeyConstraint(["id"], ["collection.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "dataset_entry",
        sa.Column("dataset_id", sa.Integer(), nullable=False),
        sa.Column("molecule_id", sa.Integer(), nullable=False),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("comment", sa.String(), nullable=True),
        sa.Column("local_results", sa.JSON(), nullable=True),
        sa.ForeignKeyConstraint(["dataset_id"], ["dataset.id"], ondelete="cascade"),
        sa.ForeignKeyConstraint(["molecule_id"], ["molecule.id"], ondelete="cascade"),
        sa.PrimaryKeyConstraint("dataset_id", "molecule_id"),
    )
    op.create_table(
        "reaction_dataset_entry",
        sa.Column("reaction_dataset_id", sa.Integer(), nullable=False),
        sa.Column("attributes", sa.JSON(), nullable=True),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("reaction_results", sa.JSON(), nullable=True),
        sa.Column("stoichiometry", sa.JSON(), nullable=True),
        sa.Column("extras", sa.JSON(), nullable=True),
        sa.ForeignKeyConstraint(["reaction_dataset_id"], ["reaction_dataset.id"], ondelete="cascade"),
        sa.PrimaryKeyConstraint("reaction_dataset_id", "name"),
    )
    op.add_column("collection", sa.Column("collection_type", sa.String(), nullable=True))
    op.add_column("collection", sa.Column("provenance", sa.JSON(), nullable=True))
    op.create_index("ix_collection_type", "collection", ["collection_type"], unique=False)
    ### end Alembic commands ###

    # ------------ copy data
    print("Start Data migration...")
    migrate_collections()


def migrate_collections():
    session = Session(bind=op.get_bind())

    collection_map = {"dataset": DatasetORM, "reactiondataset": ReactionDatasetORM}

    collections = session.query(CollectionORM)

    for collection in collections:
        print(f"collection: id:{collection.id}, lname: {collection.lname}")

        if collection.collection in collection_map.keys():
            collection_class = collection_map[collection.collection]
        else:
            continue

        fields = collection.to_dict(exclude=["id"])
        session.query(CollectionORM).filter_by(id=collection.id).delete(synchronize_session=False)
        session.commit()

        dataset = collection_class(**fields)
        session.add(dataset)
        session.commit()

        dataset.update_relations(**fields)  # will use related attr
        session.commit()

        # print(f'dataset.id: {dataset.id},\n dataset.records: {len(dataset.records)}')


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("ix_collection_type", table_name="collection")
    op.drop_column("collection", "collection_type")
    op.drop_column("collection", "provenance")
    op.drop_table("reaction_dataset_entry")
    op.drop_table("dataset_entry")
    op.drop_table("reaction_dataset")
    op.drop_table("dataset")
    # ### end Alembic commands ###
