"""Persist explicit unknown point-in-time sector membership snapshots."""

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

revision: str = "0006_membership_unknown"
down_revision: str | None = "0005_radar_daily_aggregate"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    """Allow one null-stock marker for an explicitly empty sector partition."""

    op.add_column(
        "sector_radar_membership",
        sa.Column("membership_key", sa.String(32), nullable=True),
    )
    op.execute("UPDATE sector_radar_membership SET membership_key = stock_code")
    op.drop_constraint(
        "sector_radar_membership_pkey",
        "sector_radar_membership",
        type_="primary",
    )
    op.drop_constraint(
        "ck_sector_radar_membership_status",
        "sector_radar_membership",
        type_="check",
    )
    op.alter_column(
        "sector_radar_membership",
        "membership_key",
        existing_type=sa.String(32),
        nullable=False,
    )
    op.alter_column(
        "sector_radar_membership",
        "stock_code",
        existing_type=sa.String(12),
        nullable=True,
    )
    op.alter_column(
        "sector_radar_membership",
        "stock_name",
        existing_type=sa.String(128),
        nullable=True,
    )
    op.create_primary_key(
        "sector_radar_membership_pkey",
        "sector_radar_membership",
        ["source_snapshot_id", "sector_code", "membership_key"],
    )
    op.create_check_constraint(
        "ck_sector_radar_membership_status",
        "sector_radar_membership",
        "(membership_status = 'available' "
        "AND stock_code IS NOT NULL AND stock_name IS NOT NULL "
        "AND membership_key = stock_code) OR "
        "(membership_status = 'membership_unknown' "
        "AND stock_code IS NULL AND stock_name IS NULL "
        "AND membership_key = '__membership_unknown__')",
    )


def downgrade() -> None:
    """Discard unknown markers and restore the available-member-only schema."""

    op.execute("DELETE FROM sector_radar_membership WHERE membership_status = 'membership_unknown'")
    op.drop_constraint(
        "ck_sector_radar_membership_status",
        "sector_radar_membership",
        type_="check",
    )
    op.drop_constraint(
        "sector_radar_membership_pkey",
        "sector_radar_membership",
        type_="primary",
    )
    op.alter_column(
        "sector_radar_membership",
        "stock_code",
        existing_type=sa.String(12),
        nullable=False,
    )
    op.alter_column(
        "sector_radar_membership",
        "stock_name",
        existing_type=sa.String(128),
        nullable=False,
    )
    op.drop_column("sector_radar_membership", "membership_key")
    op.create_primary_key(
        "sector_radar_membership_pkey",
        "sector_radar_membership",
        ["source_snapshot_id", "sector_code", "stock_code"],
    )
    op.create_check_constraint(
        "ck_sector_radar_membership_status",
        "sector_radar_membership",
        "membership_status = 'available'",
    )