"""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 ###
