"""Align usage_ledger with the UsageLedger ORM.

001 created usage_ledger without project_id, unit_label, metadata_json, or
reported_to_stripe. Cloud usage finalize/release SELECT those columns and
500 with UndefinedColumnError until this runs.

Revision ID: 120
Revises: 119
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect
from sqlalchemy.dialects import postgresql

revision = "120"
down_revision = "119"
branch_labels = None
depends_on = None


def upgrade() -> None:
    conn = op.get_bind()
    insp = inspect(conn)
    cols = {c["name"] for c in insp.get_columns("usage_ledger")}

    if "project_id" not in cols:
        op.add_column(
            "usage_ledger",
            sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True),
        )
        op.create_foreign_key(
            "fk_usage_ledger_project_id",
            "usage_ledger",
            "projects",
            ["project_id"],
            ["id"],
        )
        op.create_index("idx_usage_ledger_project", "usage_ledger", ["project_id"])

    if "unit_label" not in cols:
        op.add_column(
            "usage_ledger",
            sa.Column("unit_label", sa.String(length=30), nullable=True),
        )

    if "reported_to_stripe" not in cols:
        op.add_column(
            "usage_ledger",
            sa.Column(
                "reported_to_stripe",
                sa.Boolean(),
                nullable=False,
                server_default=sa.text("false"),
            ),
        )

    cols = {c["name"] for c in inspect(conn).get_columns("usage_ledger")}
    if "metadata_json" not in cols:
        if "metadata" in cols:
            op.alter_column(
                "usage_ledger",
                "metadata",
                new_column_name="metadata_json",
            )
        else:
            op.add_column(
                "usage_ledger",
                sa.Column("metadata_json", postgresql.JSONB(), nullable=True),
            )


def downgrade() -> None:
    conn = op.get_bind()
    insp = inspect(conn)
    cols = {c["name"] for c in insp.get_columns("usage_ledger")}
    indexes = {i["name"] for i in insp.get_indexes("usage_ledger")}
    fks = {fk["name"] for fk in insp.get_foreign_keys("usage_ledger")}

    if "idx_usage_ledger_project" in indexes:
        op.drop_index("idx_usage_ledger_project", table_name="usage_ledger")
    if "fk_usage_ledger_project_id" in fks:
        op.drop_constraint("fk_usage_ledger_project_id", "usage_ledger", type_="foreignkey")
    if "project_id" in cols:
        op.drop_column("usage_ledger", "project_id")
    if "unit_label" in cols:
        op.drop_column("usage_ledger", "unit_label")
    if "reported_to_stripe" in cols:
        op.drop_column("usage_ledger", "reported_to_stripe")
    if "metadata_json" in cols and "metadata" not in cols:
        op.alter_column(
            "usage_ledger",
            "metadata_json",
            new_column_name="metadata",
        )