"""Peewee migrations -- 155_MAIN.py.

Some examples (model - class or model name)::

    > Model = migrator.orm['table_name']            # Return model in current state by name
    > Model = migrator.ModelClass                   # Return model in current state by name

    > migrator.sql(sql)                             # Run custom SQL
    > migrator.run(func, *args, **kwargs)           # Run python function with the given args
    > migrator.create_model(Model)                  # Create a model (could be used as decorator)
    > migrator.remove_model(model, cascade=True)    # Remove a model
    > migrator.add_fields(model, **fields)          # Add fields to a model
    > migrator.change_fields(model, **fields)       # Change fields
    > migrator.remove_fields(model, *field_names, cascade=True)
    > migrator.rename_field(model, old_field_name, new_field_name)
    > migrator.rename_table(model, new_table_name)
    > migrator.add_index(model, *col_names, unique=False)
    > migrator.add_not_null(model, *field_names)
    > migrator.add_default(model, field_name, default)
    > migrator.add_constraint(model, name, sql)
    > migrator.drop_index(model, *col_names)
    > migrator.drop_not_null(model, *field_names)
    > migrator.drop_constraints(model, *constraints)

"""

import peewee as pw
from peewee_migrate import Migrator


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

    for car in Car.filter(Car.driver.is_null(False)).iterator():
        CarDriver.create(driver=car.driver, vehicle=car, t=0)


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

    for pump in Pump.filter(Pump.driver.is_null(False)).iterator():
        PumpDriver.create(driver=pump.driver, vehicle=pump, t=0)


def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
    """Write your migrations here."""
    
    migrator.rename_field('delivery', 'car_driver', 'driver_name')
    
    migrator.rename_field('delivery', 'car_driver_contact', 'driver_contact')

    migrator.add_fields(
        'delivery',

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

    migrator.sql("""
        UPDATE delivery 
        SET driver_record_id = car.driver_id
        from car
        WHERE delivery.car_record_id IS NOT NULL
          AND car.driver_id IS NOT NULL
          AND delivery.car_record_id = car.id
        ;
    """)

    @migrator.create_model
    class CarDriver(pw.Model):
        id = pw.AutoField()
        driver = pw.ForeignKeyField(column_name='driver_id', field='id', model=migrator.orm['driver'], on_delete='CASCADE')
        t = pw.IntegerField()
        vehicle = pw.ForeignKeyField(column_name='vehicle_id', field='id', model=migrator.orm['car'], on_delete='CASCADE')

        class Meta:
            table_name = "cardriver"
            indexes = [(('vehicle', 'driver'), True)]

    migrator.run(fill_car_driver, migrator)

    migrator.rename_field('pumporder', 'pump_driver', 'driver_name')

    migrator.rename_field('pumporder', 'pump_driver_contact', 'driver_contact')

    migrator.add_fields(
        'pumporder',

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

    migrator.sql("""
        UPDATE pumporder 
        SET driver_record_id = pump.driver_id
        from pump
        WHERE pumporder.pump_record_id IS NOT NULL
          AND pump.driver_id IS NOT NULL
          AND pumporder.pump_record_id = pump.id
        ;
    """)

    @migrator.create_model
    class PumpDriver(pw.Model):
        id = pw.AutoField()
        driver = pw.ForeignKeyField(column_name='driver_id', field='id', model=migrator.orm['driver'], on_delete='CASCADE')
        t = pw.IntegerField()
        vehicle = pw.ForeignKeyField(column_name='vehicle_id', field='id', model=migrator.orm['pump'], on_delete='CASCADE')

        class Meta:
            table_name = "pumpdriver"
            indexes = [(('vehicle', 'driver'), True)]

    migrator.run(fill_pump_driver, migrator)


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

    migrator.remove_model('cardriver')

    migrator.remove_fields('delivery', 'driver_record')
    
    migrator.rename_field('delivery', 'driver_name', 'car_driver')
    
    migrator.rename_field('delivery', 'driver_contact', 'car_driver_contact')

    migrator.remove_model('pumpdriver')

    migrator.remove_fields('pumporder', 'driver_record')

    migrator.rename_field('pumporder', 'driver_name', 'pump_driver')

    migrator.rename_field('pumporder', 'driver_contact', 'pump_driver_contact')
