"""add pension tables

Revision ID: 0007
Revises: 0006
Create Date: 2026-04-27
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision = "0007"
down_revision = "0006"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ------------------------------------------------------------------
    # pension_metadata — one row per pension account
    # ------------------------------------------------------------------
    op.create_table(
        "pension_metadata",
        sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column("user_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
        sa.Column("account_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("accounts.id", ondelete="CASCADE"), nullable=False),
        sa.Column("pension_type", sa.String(20), nullable=False),
        sa.Column("provider_name", sa.LargeBinary, nullable=True),
        sa.Column("scheme_name", sa.LargeBinary, nullable=True),
        sa.Column("member_reference", sa.LargeBinary, nullable=True),
        sa.Column("dob", sa.Date, nullable=True),
        sa.Column("target_retirement_age", sa.Integer, nullable=True),
        sa.Column("assumed_growth_rate", sa.Numeric(5, 4), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
    )
    op.create_unique_constraint("uq_pension_metadata_account_id", "pension_metadata", ["account_id"])
    op.create_index("ix_pension_metadata_user_id", "pension_metadata", ["user_id"])
    op.create_index("ix_pension_metadata_account_id", "pension_metadata", ["account_id"])

    # ------------------------------------------------------------------
    # pension_contributions — contribution history per pension
    # ------------------------------------------------------------------
    op.create_table(
        "pension_contributions",
        sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column("user_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
        sa.Column("pension_id", postgresql.UUID(as_uuid=True),
                  sa.ForeignKey("pension_metadata.id", ondelete="CASCADE"), nullable=False),
        sa.Column("contribution_date", sa.Date, nullable=False),
        sa.Column("tax_year", sa.Integer, nullable=False),
        sa.Column("member_amount", sa.Numeric(14, 2), nullable=False),
        sa.Column("employer_amount", sa.Numeric(14, 2), nullable=False, server_default="0"),
        sa.Column("relief_type", sa.String(20), nullable=False),
        sa.Column("gross_amount", sa.Numeric(14, 2), nullable=False),
        sa.Column("relief_amount", sa.Numeric(14, 2), nullable=False, server_default="0"),
        sa.Column("notes", sa.LargeBinary, nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
    )
    op.create_index("ix_pension_contributions_user_id", "pension_contributions", ["user_id"])
    op.create_index("ix_pension_contributions_pension_id", "pension_contributions", ["pension_id"])
    op.create_index("ix_pension_contributions_date", "pension_contributions", ["contribution_date"])
    op.create_index("ix_pension_contributions_tax_year", "pension_contributions", ["tax_year"])

    # ------------------------------------------------------------------
    # RLS
    # ------------------------------------------------------------------
    for table in ["pension_metadata", "pension_contributions"]:
        op.execute(f"ALTER TABLE {table} ENABLE ROW LEVEL SECURITY")
        op.execute(f"""
            CREATE POLICY {table}_user_isolation ON {table}
            USING (user_id = current_app_user_id())
        """)


def downgrade() -> None:
    for table in ["pension_metadata", "pension_contributions"]:
        op.execute(f"DROP POLICY IF EXISTS {table}_user_isolation ON {table}")
        op.execute(f"ALTER TABLE {table} DISABLE ROW LEVEL SECURITY")

    op.drop_table("pension_contributions")
    op.drop_table("pension_metadata")