"""Peewee migrations -- 129_MAIN.py.

Continues in 130_MAIN.py
"""

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_fields(
        'delivery',
        _search_string=pw.TextField(null=True),
    )

    migrator.add_fields(
        'order',
        continuous_mode=pw.BooleanField(default=False),
    )

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

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

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

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

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


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

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