"""Peewee migrations -- 130_MAIN.py.

Continuation of 129_MAIN
"""

from contextlib import suppress

import peewee as pw
from peewee_migrate import Migrator

from migrations.utils import drop_fk_column_fix


with suppress(ImportError):
    import playhouse.postgres_ext as pw_pext


drop_fk_column_fix()


def migrate(migrator: Migrator, database: pw.Database, *, fake: bool = False) -> None:
    """Write your migrations here."""

    migrator.add_not_null('batch', 'delivery')


def rollback(migrator: Migrator, database: pw.Database, *, fake: bool = False) -> None:
    """Write your rollback migrations here."""

    migrator.remove_fields('delivery', '_search_string')

    migrator.remove_fields('order', 'continuous_mode')

    # Re-link Batch from Delivery to Order
    migrator.add_fields(
        'batch',
        delivery=pw.ForeignKeyField(column_name='order_id', null=True, field='id', model=migrator.orm['order'], on_delete='CASCADE')
    )

    migrator.sql("""
        UPDATE batch
        SET order_id = delivery.order_id
        FROM (SELECT id, order_id FROM delivery) AS delivery
        WHERE batch.delivery_id = delivery.id;
    """)

    migrator.add_index('batch', 'order', 'batch_number', unique=True)

    # Remove Batch.delivery
    migrator.drop_index('batch', 'delivery', 'batch_number')

    migrator.remove_fields('batch', 'delivery')
