"""Add country_iso to customer_addresses

Revision ID: r6f7a8b9c0d1
Revises: q5e6f7a8b9c0
Create Date: 2026-01-02

Adds country_iso field to customer_addresses table and renames
country to country_name for clarity.

This migration is idempotent - it checks for existing columns before
making changes.
"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy import text


# revision identifiers, used by Alembic.
revision = "r6f7a8b9c0d1"
down_revision = "q5e6f7a8b9c0"
branch_labels = None
depends_on = None


# Country name to ISO code mapping for backfill
COUNTRY_ISO_MAP = {
    "Luxembourg": "LU",
    "Germany": "DE",
    "France": "FR",
    "Belgium": "BE",
    "Netherlands": "NL",
    "Austria": "AT",
    "Italy": "IT",
    "Spain": "ES",
    "Portugal": "PT",
    "Poland": "PL",
    "Czech Republic": "CZ",
    "Czechia": "CZ",
    "Slovakia": "SK",
    "Hungary": "HU",
    "Romania": "RO",
    "Bulgaria": "BG",
    "Greece": "GR",
    "Croatia": "HR",
    "Slovenia": "SI",
    "Estonia": "EE",
    "Latvia": "LV",
    "Lithuania": "LT",
    "Finland": "FI",
    "Sweden": "SE",
    "Denmark": "DK",
    "Ireland": "IE",
    "Cyprus": "CY",
    "Malta": "MT",
    "United Kingdom": "GB",
    "Switzerland": "CH",
    "United States": "US",
}


def get_column_names(connection, table_name):
    """Get list of column names for a table (PostgreSQL)."""
    result = connection.execute(text(
        "SELECT column_name FROM information_schema.columns "
        "WHERE table_name = :table AND table_schema = 'public'"
    ), {"table": table_name})
    return [row[0] for row in result]


def upgrade() -> None:
    connection = op.get_bind()
    columns = get_column_names(connection, "customer_addresses")

    # Check if we need to do anything (idempotent check)
    has_country = "country" in columns
    has_country_name = "country_name" in columns
    has_country_iso = "country_iso" in columns

    # If already has new columns, nothing to do
    if has_country_name and has_country_iso:
        print("  Columns country_name and country_iso already exist, skipping")
        return

    # If has old 'country' column, rename it (PostgreSQL supports direct rename)
    if has_country and not has_country_name:
        op.alter_column(
            "customer_addresses",
            "country",
            new_column_name="country_name",
        )

    # Add country_iso if it doesn't exist
    if not has_country_iso:
        op.add_column(
            "customer_addresses",
            sa.Column("country_iso", sa.String(5), nullable=True)
        )

        # Backfill country_iso from country_name
        for country_name, iso_code in COUNTRY_ISO_MAP.items():
            connection.execute(
                text(
                    "UPDATE customer_addresses SET country_iso = :iso "
                    "WHERE country_name = :name"
                ),
                {"iso": iso_code, "name": country_name},
            )

        # Set default for any remaining NULL values
        connection.execute(
            text(
                "UPDATE customer_addresses SET country_iso = 'LU' "
                "WHERE country_iso IS NULL"
            )
        )

        # Make country_iso NOT NULL (PostgreSQL supports direct alter)
        op.alter_column(
            "customer_addresses",
            "country_iso",
            existing_type=sa.String(5),
            nullable=False,
        )


def downgrade() -> None:
    connection = op.get_bind()
    columns = get_column_names(connection, "customer_addresses")

    has_country_name = "country_name" in columns
    has_country_iso = "country_iso" in columns
    has_country = "country" in columns

    # Only downgrade if in the new state
    if has_country_name and not has_country:
        op.alter_column(
            "customer_addresses",
            "country_name",
            new_column_name="country",
        )

    if has_country_iso:
        op.drop_column("customer_addresses", "country_iso")