"""Peewee migrations -- 138_MAIN.py and 139_MAIN.py

Split into two due to field removal being stupid
"""

from contextlib import suppress

import peewee as pw
from peewee_migrate import Migrator


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


RECIPE_PRICE_NOTE = 'Using price directly from recipe'


def fix_order_data(migrator: Migrator) -> None:
    for order in migrator.orm["order"].select():
        if order.price_concrete_modified:
            # The price was set manually, but for the whole order
            order.best_price = order.price_concrete_modified / (order.volume or 1)
            order.best_price_modifier = f"={order.best_price:.2f}"

        if order.best_price is not None:
            if order.best_price_note == RECIPE_PRICE_NOTE:
                # The source of the best price was the recipe, so we can set that value
                order.r_price = order.best_price
            else:
                # The price was set by Price.calc_best_price
                order.best_price_modifier = f"={order.best_price:.2f}"

        order.save(only=["r_price", "best_price", "best_price_modifier"])


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

    migrator.rename_field("order", "r_price", "best_price")

    migrator.add_fields(
        "order",

        r_price=pw.DoubleField(null=True),
        best_price_modifier=pw.TextField(null=True))

    migrator.run(fix_order_data, migrator)


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

    migrator.remove_fields("order", "r_price", "best_price_modifier")

    migrator.rename_field("order", "best_price", "r_price")

    migrator.rename_field("order", "best_price_note", "r_price_note")
