"""Add missing columns to fix database schema

Revision ID: 006
Revises: 005
Create Date: 2025-11-17 05:35:00.000000

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '006'
down_revision = '005'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Добавляем отсутствующий столбец account_id в participations (если еще не существует)
    op.execute("""
        DO $$ 
        BEGIN 
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='participations' AND column_name='account_id') THEN
                ALTER TABLE participations ADD COLUMN account_id INTEGER;
                ALTER TABLE participations 
                ADD CONSTRAINT fk_participations_account_id 
                FOREIGN KEY (account_id) REFERENCES accounts(id) 
                ON DELETE SET NULL;
            END IF;
        END $$;
    """)
    
    # Добавляем отсутствующие столбцы в winners
    op.execute("""
        DO $$ 
        BEGIN 
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='winners' AND column_name='is_notified') THEN
                ALTER TABLE winners ADD COLUMN is_notified BOOLEAN DEFAULT FALSE;
            END IF;
            
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='winners' AND column_name='is_claimed') THEN
                ALTER TABLE winners ADD COLUMN is_claimed BOOLEAN DEFAULT FALSE;
            END IF;
            
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='winners' AND column_name='claimed_at') THEN
                ALTER TABLE winners ADD COLUMN claimed_at TIMESTAMP WITH TIME ZONE;
            END IF;
        END $$;
    """)


def downgrade() -> None:
    # Удаляем добавленные столбцы в обратном порядке
    
    # Удаляем столбцы из winners
    op.execute("""
        DO $$ 
        BEGIN 
            IF EXISTS (SELECT 1 FROM information_schema.columns 
                      WHERE table_name='winners' AND column_name='claimed_at') THEN
                ALTER TABLE winners DROP COLUMN claimed_at;
            END IF;
            
            IF EXISTS (SELECT 1 FROM information_schema.columns 
                      WHERE table_name='winners' AND column_name='is_claimed') THEN
                ALTER TABLE winners DROP COLUMN is_claimed;
            END IF;
            
            IF EXISTS (SELECT 1 FROM information_schema.columns 
                      WHERE table_name='winners' AND column_name='is_notified') THEN
                ALTER TABLE winners DROP COLUMN is_notified;
            END IF;
        END $$;
    """)
    
    # Удаляем account_id из participations
    op.execute("""
        DO $$ 
        BEGIN 
            IF EXISTS (SELECT 1 FROM information_schema.columns 
                      WHERE table_name='participations' AND column_name='account_id') THEN
                ALTER TABLE participations DROP CONSTRAINT IF EXISTS fk_participations_account_id;
                ALTER TABLE participations DROP COLUMN account_id;
            END IF;
        END $$;
    """)