"""Peewee migrations -- 182_MAIN.py.

Introduce ``TaxDocument`` and migrate ``Order.invoice_number`` / ``Order.payment_type``
into it.

Steps:
1. Create ``taxdocument`` table (``customer_record_id`` is nullable — Czech VAT
   rules allow receipts under the small-amount threshold to omit the customer).
2. Add ``order.taxdocument_id`` nullable FK.
3. Backfill: group orders by (invoice_number, customer_record_id, payment_type),
   create one TaxDocument per group, link the orders. Orphan orders (invoice
   number but no customer) get a TaxDocument with ``customer_record_id=NULL``.
4. Drop ``order.invoice_number`` and ``order.payment_type``.
"""

import logging

import peewee as pw
from peewee_migrate import Migrator

from migrations.utils import drop_fk_column_fix


def backfill_taxdocuments(migrator: Migrator) -> None:
    """Create one TaxDocument per (invoice_number, customer_record, payment_type) group
    and link the orders. Uses raw SQL because by the time this op runs, the migrator
    has already removed invoice_number / payment_type from the Order ORM model
    (the remove_fields call mutates the in-memory model eagerly, even though the
    actual DROP COLUMN runs after this).
    """
    database = migrator.__database__

    # Read setup defaults
    setup_rows = list(database.execute_sql("SELECT vat_rate, currency_symbol FROM setup LIMIT 1"))
    default_vat_rate = float(setup_rows[0][0]) if setup_rows and setup_rows[0][0] is not None else None
    default_currency = setup_rows[0][1] if setup_rows else None

    # Group orders by (invoice_number, customer_record_id, payment_type).
    # customer_record_id may be NULL — legacy orphan invoices with no customer
    # still get a TaxDocument (with customer_record_id=NULL).
    rows = list(database.execute_sql('SELECT id, t, invoice_number, customer_record_id, payment_type FROM "order" WHERE invoice_number IS NOT NULL'))

    groups: dict[tuple, list[tuple[int, int]]] = {}
    for order_id, t_val, invoice_number, customer_id, payment_type in rows:
        key = (invoice_number, customer_id, payment_type)
        groups.setdefault(key, []).append((order_id, t_val))

    by_invoice_num: dict[int, int] = {}
    for invoice_number, _c, _p in groups:
        by_invoice_num[invoice_number] = by_invoice_num.get(invoice_number, 0) + 1

    splits = sum(1 for c in by_invoice_num.values() if c > 1)
    if splits:
        logging.warning(f"taxdocument backfill: {splits} invoice_number(s) split across (customer, payment_type) groups")

    # When invoice_number is shared across groups (split), only the first group keeps
    # the original number; the rest get fresh numbers from the high end.
    next_free = max(by_invoice_num.keys(), default=0) + 1
    used_numbers: set[int] = set()

    for (invoice_number, customer_id, payment_type), order_pairs in sorted(
        groups.items(),
        key=lambda kv: (kv[0][0], kv[0][1] is None, kv[0][1], kv[0][2] is None, kv[0][2]),
    ):
        if invoice_number in used_numbers:
            doc_number = next_free
            next_free += 1
        else:
            doc_number = invoice_number
        used_numbers.add(doc_number)

        first_t = min((p[1] for p in order_pairs if p[1] is not None), default=0)
        order_ids = [p[0] for p in order_pairs]

        database.execute_sql(
            """INSERT INTO taxdocument (t, number, customer_record_id, date, payment_type,
                vat_rate, currency, note,
                audit_created_by, audit_created_at, audit_changed_by, audit_changed_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                first_t,
                doc_number,
                customer_id,
                first_t,
                payment_type if payment_type is not None else 0,
                default_vat_rate,
                default_currency,
                "",
                "migration",
                first_t,
                "migration",
                first_t,
            ),
        )

        doc_id = list(database.execute_sql("SELECT id FROM taxdocument WHERE number = ?", (doc_number,)))[0][0]

        # Link orders to the new document
        placeholders = ",".join(["?"] * len(order_ids))
        database.execute_sql(
            f'UPDATE "order" SET taxdocument_id = ? WHERE id IN ({placeholders})',
            (doc_id, *order_ids),
        )

    logging.info(f"taxdocument backfill: created {len(groups)} documents from {len(rows)} orders ({splits} invoice_number splits)")


def migrate(migrator: Migrator, database: pw.Database, *, fake: bool = False) -> None:
    @migrator.create_model
    class TaxDocument(pw.Model):
        id = pw.AutoField()
        audit_created_by = pw.TextField(null=True)
        audit_created_at = pw.IntegerField(constraints=[pw.SQL("DEFAULT 0")], default=0)
        audit_changed_by = pw.TextField(null=True)
        audit_changed_at = pw.IntegerField(constraints=[pw.SQL("DEFAULT 0")], default=0)
        t = pw.IntegerField()
        number = pw.IntegerField(unique=True)
        customer_record = pw.ForeignKeyField(
            backref="taxdocuments",
            column_name="customer_record_id",
            field="id",
            model=migrator.orm["customer"],
            null=True,
            on_delete="RESTRICT",
        )
        date = pw.IntegerField()
        payment_type = pw.IntegerField()
        vat_rate = pw.DoubleField(null=True)
        currency = pw.TextField(null=True)
        note = pw.TextField(null=True, default="")

        class Meta:
            table_name = "taxdocument"

    migrator.add_fields(
        "order",
        taxdocument=pw.ForeignKeyField(
            backref="orders",
            column_name="taxdocument_id",
            field="id",
            model=migrator.orm["taxdocument"],
            null=True,
            on_delete="SET NULL",
        ),
    )

    migrator.run(backfill_taxdocuments, migrator)

    drop_fk_column_fix()
    migrator.remove_fields("order", "invoice_number", "payment_type")


def rollback(migrator: Migrator, database: pw.Database, *, fake: bool = False) -> None:
    migrator.add_fields(
        "order",
        invoice_number=pw.IntegerField(null=True),
        payment_type=pw.IntegerField(null=True),
    )

    drop_fk_column_fix()
    migrator.remove_fields("order", "taxdocument")
    migrator.remove_model("taxdocument")
