"""Peewee migrations -- 177_MAIN.py.

Add Pricelist and PricelistRecipe models.
Add pricelist FK to Customer and Order.
Migrate existing Recipe.price data into PricelistRecipe(pricelist=NULL).
Drop Recipe.price column.
"""

import peewee as pw
from peewee_migrate import Migrator

from migrations.utils import drop_fk_column_fix


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

    # 1. Create Pricelist table
    @migrator.create_model
    class Pricelist(pw.Model):
        id = pw.AutoField()
        hidden = pw.BooleanField(default=False)
        _search_string = pw.TextField(null=True)
        name = pw.TextField()

        class Meta:
            table_name = "pricelist"

    # 2. Create PricelistRecipe table
    @migrator.create_model
    class PricelistRecipe(pw.Model):
        id = pw.AutoField()
        _search_string = pw.TextField(null=True)
        pricelist = pw.ForeignKeyField(column_name='pricelist_id', field='id', model=migrator.orm['pricelist'], null=True, on_delete='CASCADE')
        recipe = pw.ForeignKeyField(column_name='recipe_id', field='id', model=migrator.orm['recipe'], on_delete='CASCADE')
        price = pw.DoubleField()

        class Meta:
            table_name = "pricelistrecipe"
            indexes = [(('pricelist', 'recipe'), True)]

    # 3. Add pricelist FK to Customer
    migrator.add_fields(
        'customer',
        pricelist=pw.ForeignKeyField(column_name='pricelist_id', field='id', model=migrator.orm['pricelist'], null=True, on_delete='SET NULL'))

    # 4. Add pricelist FK to Order
    migrator.add_fields(
        'order',
        pricelist=pw.ForeignKeyField(column_name='pricelist_id', field='id', model=migrator.orm['pricelist'], null=True, on_delete='SET NULL'))

    # 5. Migrate existing Recipe.price → PricelistRecipe(pricelist=NULL)
    def migrate_prices(database):
        database.execute_sql(
            "INSERT INTO pricelistrecipe (_search_string, pricelist_id, recipe_id, price) "
            "SELECT NULL, NULL, id, price FROM recipe WHERE price IS NOT NULL"
        )

    migrator.run(migrate_prices, database)

    # 6. Drop Recipe.price column
    drop_fk_column_fix()
    migrator.remove_fields('recipe', 'price')


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

    # Re-add Recipe.price
    migrator.add_fields(
        'recipe',
        price=pw.DoubleField(null=True))

    # Restore prices from PricelistRecipe defaults back to Recipe
    def restore_prices(database):
        database.execute_sql(
            "UPDATE recipe SET price = ("
            "  SELECT price FROM pricelistrecipe "
            "  WHERE pricelistrecipe.recipe_id = recipe.id AND pricelistrecipe.pricelist_id IS NULL"
            ")"
        )

    migrator.run(restore_prices, database)

    drop_fk_column_fix()
    migrator.remove_fields('order', 'pricelist')

    drop_fk_column_fix()
    migrator.remove_fields('customer', 'pricelist')

    migrator.remove_model('pricelistrecipe')
    migrator.remove_model('pricelist')
