"""investor access: username, entity access, documents, capital accounts

Revision ID: a1b2c3d4e5f6
Revises: 2792c4ff4612
Create Date: 2026-06-26 14:30:00.000000

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlmodel


# revision identifiers, used by Alembic.
revision: str = 'a1b2c3d4e5f6'
down_revision: Union[str, None] = '2792c4ff4612'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    # users: add username (login handle), make email optional.
    with op.batch_alter_table('users', schema=None) as batch_op:
        batch_op.add_column(
            sa.Column('username', sqlmodel.sql.sqltypes.AutoString(), nullable=True)
        )
        batch_op.alter_column('email', existing_type=sa.String(), nullable=True)

    # Backfill username for any existing internal accounts so the NOT NULL holds.
    op.execute("UPDATE users SET username = email WHERE username IS NULL")

    with op.batch_alter_table('users', schema=None) as batch_op:
        batch_op.alter_column(
            'username',
            existing_type=sqlmodel.sql.sqltypes.AutoString(),
            nullable=False,
        )
        batch_op.create_unique_constraint('uq_users_username', ['username'])

    op.create_table(
        'entity_access',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('user_id', sa.Integer(), nullable=False),
        sa.Column('entity_id', sa.Integer(), nullable=False),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ),
        sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('user_id', 'entity_id', name='uq_access_user_entity'),
    )

    op.create_table(
        'documents',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('entity_id', sa.Integer(), nullable=False),
        sa.Column('investor_user_id', sa.Integer(), nullable=True),
        sa.Column('category', sa.Enum('capital_account', 'k1', 'statement', 'tax', 'other', name='documentcategory'), nullable=False),
        sa.Column('title', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('original_filename', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('content_type', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('size_bytes', sa.Integer(), nullable=False),
        sa.Column('storage_path', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('uploaded_by', sa.Integer(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ),
        sa.ForeignKeyConstraint(['investor_user_id'], ['users.id'], ),
        sa.ForeignKeyConstraint(['uploaded_by'], ['users.id'], ),
        sa.PrimaryKeyConstraint('id'),
    )

    op.create_table(
        'capital_account_statements',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('entity_id', sa.Integer(), nullable=False),
        sa.Column('investor_user_id', sa.Integer(), nullable=False),
        sa.Column('as_of_date', sa.Date(), nullable=False),
        sa.Column('beginning_balance_cents', sa.Integer(), nullable=False),
        sa.Column('contributions_cents', sa.Integer(), nullable=False),
        sa.Column('distributions_cents', sa.Integer(), nullable=False),
        sa.Column('ending_balance_cents', sa.Integer(), nullable=False),
        sa.Column('document_id', sa.Integer(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(['document_id'], ['documents.id'], ),
        sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ),
        sa.ForeignKeyConstraint(['investor_user_id'], ['users.id'], ),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('entity_id', 'investor_user_id', 'as_of_date', name='uq_capacct_entity_investor_date'),
    )


def downgrade() -> None:
    """Downgrade schema."""
    op.drop_table('capital_account_statements')
    op.drop_table('documents')
    op.drop_table('entity_access')
    with op.batch_alter_table('users', schema=None) as batch_op:
        batch_op.drop_constraint('uq_users_username', type_='unique')
        batch_op.alter_column('email', existing_type=sa.String(), nullable=False)
        batch_op.drop_column('username')