"""Add signup management fields to accounts table

Revision ID: 0005
Revises: 0004
Create Date: 2026-06-24 00:00:00.000000
"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

revision: str = "0005"
down_revision: Union[str, None] = "0004a"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Add new columns to accounts table
    op.add_column(
        "accounts",
        sa.Column("credentials_sent_at", sa.DateTime(timezone=True), nullable=True),
    )
    op.add_column(
        "accounts",
        sa.Column("account_manager", sa.String(100), nullable=True),
    )
    op.add_column(
        "accounts",
        sa.Column("size", sa.String(50), nullable=True),
    )
    op.add_column(
        "accounts",
        sa.Column(
            "signup_source",
            sa.String(20),
            nullable=False,
            server_default="marketing",
        ),
    )
    op.add_column(
        "accounts",
        sa.Column("activated_at", sa.DateTime(timezone=True), nullable=True),
    )
    op.add_column(
        "accounts",
        sa.Column(
            "feature_flags",
            postgresql.JSONB(astext_type=sa.Text()),
            nullable=True,
            server_default="{}",
        ),
    )

    # Backfill size from config_json where present
    op.execute(
        """
        UPDATE accounts
        SET size = config_json->>'size'
        WHERE config_json IS NOT NULL
          AND config_json->>'size' IS NOT NULL
          AND size IS NULL
        """
    )

    # Backfill signup_source from config_json where present
    op.execute(
        """
        UPDATE accounts
        SET signup_source = COALESCE(config_json->>'signup_source', 'marketing')
        WHERE config_json IS NOT NULL
        """
    )

    # Create indexes for common filter columns
    op.create_index("ix_accounts_signup_source", "accounts", ["signup_source"])
    op.create_index("ix_accounts_credentials_sent_at", "accounts", ["credentials_sent_at"])


def downgrade() -> None:
    op.drop_index("ix_accounts_credentials_sent_at", table_name="accounts")
    op.drop_index("ix_accounts_signup_source", table_name="accounts")
    op.drop_column("accounts", "feature_flags")
    op.drop_column("accounts", "activated_at")
    op.drop_column("accounts", "signup_source")
    op.drop_column("accounts", "size")
    op.drop_column("accounts", "account_manager")
    op.drop_column("accounts", "credentials_sent_at")