"""Peewee migrations -- Continuation of 155_MAIN.py because the remove_fields is retarded
"""

import peewee as pw
from peewee_migrate import Migrator


def spread_car_driver(migrator: Migrator) -> None:
    CarDriver = migrator.orm["cardriver"]

    for car_driver in CarDriver.select().iterator():
        car_driver.vehicle.driver = car_driver.driver
        car_driver.vehicle.save(only=["driver"])


def spread_pump_driver(migrator: Migrator) -> None:
    PumpDriver = migrator.orm["pumpdriver"]

    for pump_driver in PumpDriver.select().iterator():
        pump_driver.vehicle.driver = pump_driver.driver
        pump_driver.vehicle.save(only=["driver"])


def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
    """Write your migrations here."""
    
    migrator.remove_fields('car', 'driver')

    migrator.remove_fields('pump', 'driver')


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

    migrator.add_fields(
        'car',

        driver=pw.ForeignKeyField(column_name='driver_id', field='id', model=migrator.orm['driver'], null=True, on_delete='SET NULL'))

    migrator.run(spread_car_driver, migrator)

    migrator.add_fields(
        'pump',

        driver=pw.ForeignKeyField(column_name='driver_id', field='id', model=migrator.orm['driver'], null=True, on_delete='SET NULL'))
            
    migrator.run(spread_pump_driver, migrator)
