"""add recurring expenses

Revision ID: add_recurring_expenses
Revises: # You'll need to update this with your latest migration
Create Date: 2024-03-19 10:00:00.000000

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

# revision identifiers, used by Alembic.
revision = 'add_recurring_expenses'
down_revision = None  # Update this with your latest migration
branch_labels = None
depends_on = None

def upgrade() -> None:
    # Create recurrence_patterns table
    op.create_table(
        'recurrence_patterns',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('type', sa.String(), nullable=False),
        sa.Column('interval', sa.Integer(), nullable=False),
        sa.Column('days_of_week', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.Column('end_date', sa.DateTime(), nullable=True),
        sa.Column('max_occurrences', sa.Integer(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_recurrence_patterns_id'), 'recurrence_patterns', ['id'], unique=False)

    # Add recurring expense columns to expenses table
    op.add_column('expenses', sa.Column('is_recurring', sa.Boolean(), nullable=False, server_default='false'))
    op.add_column('expenses', sa.Column('next_occurrence', sa.DateTime(), nullable=True))
    op.add_column('expenses', sa.Column('last_occurrence', sa.DateTime(), nullable=True))
    op.add_column('expenses', sa.Column('recurrence_pattern_id', sa.Integer(), nullable=True))
    op.add_column('expenses', sa.Column('parent_expense_id', sa.Integer(), nullable=True))

    # Add foreign key constraints
    op.create_foreign_key(
        'fk_expenses_recurrence_pattern_id',
        'expenses', 'recurrence_patterns',
        ['recurrence_pattern_id'], ['id'],
        ondelete='SET NULL'
    )
    op.create_foreign_key(
        'fk_expenses_parent_expense_id',
        'expenses', 'expenses',
        ['parent_expense_id'], ['id'],
        ondelete='SET NULL'
    )

    # Add indexes
    op.create_index(
        'ix_expenses_recurring_next_occurrence',
        'expenses',
        ['is_recurring', 'next_occurrence'],
        postgresql_where=sa.text('is_recurring = true')
    )

def downgrade() -> None:
    # Drop indexes
    op.drop_index('ix_expenses_recurring_next_occurrence', table_name='expenses')

    # Drop foreign key constraints
    op.drop_constraint('fk_expenses_parent_expense_id', 'expenses', type_='foreignkey')
    op.drop_constraint('fk_expenses_recurrence_pattern_id', 'expenses', type_='foreignkey')

    # Drop columns from expenses table
    op.drop_column('expenses', 'parent_expense_id')
    op.drop_column('expenses', 'recurrence_pattern_id')
    op.drop_column('expenses', 'last_occurrence')
    op.drop_column('expenses', 'next_occurrence')
    op.drop_column('expenses', 'is_recurring')

    # Drop recurrence_patterns table
    op.drop_index(op.f('ix_recurrence_patterns_id'), table_name='recurrence_patterns')
    op.drop_table('recurrence_patterns')