"""Backfill business payment relationship columns for upgraded MySQL/TiDB databases.

Revision ID: 002
Revises: 001
Create Date: 2026-03-24
"""

import sqlalchemy as sa

from alembic import op

revision = "002"
down_revision = "001"
branch_labels = None
depends_on = None

BUSINESS_PAYMENT_TABLE = "asset_business_payments"
DEFAULT_RELATION_TYPE = sa.String(36)
REQUIRED_RELATIONS = (
    (
        "software_catalog_id",
        "asset_software_catalog",
        "fk_asset_business_payments_software_catalog_id",
    ),
    (
        "procurement_order_id",
        "asset_procurement_orders",
        "fk_asset_business_payments_procurement_order_id",
    ),
    ("supplier_id", "asset_suppliers", "fk_asset_business_payments_supplier_id"),
    (
        "requester_person_id",
        "identity_persons",
        "fk_asset_business_payments_requester_person_id",
    ),
)


def _get_inspector() -> sa.Inspector:
    return sa.inspect(op.get_bind())


def _table_exists(inspector: sa.Inspector, table_name: str) -> bool:
    return table_name in inspector.get_table_names()


def _column_names(inspector: sa.Inspector, table_name: str) -> set[str]:
    return {column["name"] for column in inspector.get_columns(table_name)}


def _columns_by_name(inspector: sa.Inspector, table_name: str) -> dict[str, dict]:
    return {
        column["name"]: column
        for column in inspector.get_columns(table_name)
    }


def _foreign_keys_by_column(
    inspector: sa.Inspector,
    table_name: str,
) -> dict[tuple[str, ...], str | None]:
    return {
        tuple(foreign_key["constrained_columns"]): foreign_key.get("name")
        for foreign_key in inspector.get_foreign_keys(table_name)
    }


def _column_type_signature(column_type: sa.types.TypeEngine) -> str:
    return str(column_type).lower().replace(" ", "")


def _target_column_type(inspector: sa.Inspector, referred_table: str) -> sa.types.TypeEngine:
    if not _table_exists(inspector, referred_table):
        return DEFAULT_RELATION_TYPE
    referred_columns = _columns_by_name(inspector, referred_table)
    referred_id = referred_columns.get("id")
    if not referred_id:
        return DEFAULT_RELATION_TYPE
    return referred_id["type"]


def upgrade() -> None:
    inspector = _get_inspector()
    if not _table_exists(inspector, BUSINESS_PAYMENT_TABLE):
        return

    existing_columns = _column_names(inspector, BUSINESS_PAYMENT_TABLE)
    missing_columns = []
    for column_name, referred_table, _ in REQUIRED_RELATIONS:
        if column_name in existing_columns:
            continue
        missing_columns.append(
            sa.Column(
                column_name,
                _target_column_type(inspector, referred_table),
                nullable=True,
            )
        )
    if missing_columns:
        with op.batch_alter_table(BUSINESS_PAYMENT_TABLE) as batch_op:
            for column in missing_columns:
                batch_op.add_column(column)

    inspector = _get_inspector()
    existing_columns = _columns_by_name(inspector, BUSINESS_PAYMENT_TABLE)
    type_updates: list[tuple[str, sa.types.TypeEngine, sa.types.TypeEngine]] = []
    for column_name, referred_table, _ in REQUIRED_RELATIONS:
        existing_column = existing_columns.get(column_name)
        if not existing_column:
            continue
        target_type = _target_column_type(inspector, referred_table)
        if _column_type_signature(existing_column["type"]) == _column_type_signature(target_type):
            continue
        type_updates.append((column_name, existing_column["type"], target_type))

    if type_updates:
        with op.batch_alter_table(BUSINESS_PAYMENT_TABLE) as batch_op:
            for column_name, existing_type, target_type in type_updates:
                batch_op.alter_column(
                    column_name,
                    existing_type=existing_type,
                    type_=target_type,
                    existing_nullable=True,
                )

    inspector = _get_inspector()
    existing_columns = _column_names(inspector, BUSINESS_PAYMENT_TABLE)
    existing_foreign_keys = _foreign_keys_by_column(inspector, BUSINESS_PAYMENT_TABLE)

    with op.batch_alter_table(BUSINESS_PAYMENT_TABLE) as batch_op:
        for column_name, referred_table, constraint_name in REQUIRED_RELATIONS:
            if column_name not in existing_columns:
                continue
            if (column_name,) in existing_foreign_keys:
                continue
            if not _table_exists(inspector, referred_table):
                continue
            batch_op.create_foreign_key(
                constraint_name,
                referred_table,
                [column_name],
                ["id"],
            )


def downgrade() -> None:
    inspector = _get_inspector()
    if not _table_exists(inspector, BUSINESS_PAYMENT_TABLE):
        return

    existing_columns = _column_names(inspector, BUSINESS_PAYMENT_TABLE)
    existing_foreign_keys = _foreign_keys_by_column(inspector, BUSINESS_PAYMENT_TABLE)

    with op.batch_alter_table(BUSINESS_PAYMENT_TABLE) as batch_op:
        for column_name, _, constraint_name in REQUIRED_RELATIONS:
            fk_name = existing_foreign_keys.get((column_name,))
            if fk_name:
                batch_op.drop_constraint(fk_name, type_="foreignkey")
            elif constraint_name in existing_foreign_keys.values():
                batch_op.drop_constraint(constraint_name, type_="foreignkey")

    inspector = _get_inspector()
    existing_columns = _column_names(inspector, BUSINESS_PAYMENT_TABLE)
    removable_columns = [
        column_name for column_name, _, _ in REQUIRED_RELATIONS if column_name in existing_columns
    ]
    if removable_columns:
        with op.batch_alter_table(BUSINESS_PAYMENT_TABLE) as batch_op:
            for column_name in removable_columns:
                batch_op.drop_column(column_name)