"""КRecreating tables min batches

Revision ID: 8ea062dfb9d0
Revises: 23130213b059
Create Date: 2024-04-08 14:38:27.571208

"""
from alembic import op
from clickhouse_sqlalchemy import engines
from clickhouse_sqlalchemy.types import common
import sqlalchemy as sa



# revision identifiers, used by Alembic.
revision = '8ea062dfb9d0'
down_revision = '23130213b059'
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()

    op.create_table(
        'tetris_min_batches_data_a',
        sa.Column("tetris_scenario_id", common.UInt32(), nullable=False),
        sa.Column('plant_id', common.UInt32(), nullable=False),
        sa.Column('material_id', common.UInt64(), nullable=False),
        sa.Column('line_id', common.String(), nullable=False),
        sa.Column('min_batch_kg', common.Float64(), nullable=False),
        sa.Column("validation_error", common.UInt8(), nullable=False),
        engines.ReplacingMergeTree(
            primary_key=['tetris_scenario_id', 'material_id', 'plant_id', 'line_id'],
            partition_by=['tetris_scenario_id'],
            order_by=['tetris_scenario_id', 'material_id', 'plant_id', 'line_id']
        )
    )

    tetris_min_batches_table = sa.sql.table(
        "tetris_min_batches_data_a",
        sa.Column("tetris_scenario_id", common.UInt32(), nullable=False),
        sa.Column('plant_id', common.UInt32(), nullable=False),
        sa.Column('material_id', common.UInt64(), nullable=False),
        sa.Column('line_id', common.String(), nullable=False),
        sa.Column('min_batch_kg', common.Float64(), nullable=False),
        sa.Column("validation_error", common.UInt8(), nullable=False),
    )



    res = conn.execute(
        "Select tetris_scenario_id, material_id, plant_id, min_batch_kg, validation_error "
        f"FROM tetris_min_batches_data"
    )
    results = res.fetchall()
    records = []
    for result in results:
        records.append(
            {
                "tetris_scenario_id": result[0],
                "material_id": result[1],
                "plant_id": result[2],
                "min_batch_kg": result[3],
                "validation_error": result[5],
            }
        )

        op.bulk_insert(tetris_min_batches_table, records)

    op.drop_table("tetris_min_batches_data")
    op.execute("rename table tetris_min_batches_data_a to tetris_min_batches_data")
    op.execute("OPTIMIZE TABLE tetris_min_batches_data FINAL")


def downgrade():
    pass