"""add watch history columns

Revision ID: 8911624d0776
Revises: 
Create Date: 2024-xx-xx xx:xx:xx.xxx
"""
from alembic import op
import sqlalchemy as sa
from datetime import datetime
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy import text


# revision identifiers, used by Alembic.
revision = '8911624d0776'
down_revision = None
branch_labels = None
depends_on = None


def column_exists(table, column):
    # Get inspector
    conn = op.get_bind()
    inspector = Inspector.from_engine(conn)
    columns = [c['name'] for c in inspector.get_columns(table)]
    return column in columns


def upgrade():
    # Add new columns if they don't exist
    with op.batch_alter_table('watch_history') as batch_op:
        # Add category column
        if not column_exists('watch_history', 'category'):
            batch_op.add_column(sa.Column('category', sa.String(100), nullable=True))
            
        # Add view_count column
        if not column_exists('watch_history', 'view_count'):
            batch_op.add_column(sa.Column('view_count', sa.Integer(), nullable=True))
            
        # Add subscriber_count column
        if not column_exists('watch_history', 'subscriber_count'):
            batch_op.add_column(sa.Column('subscriber_count', sa.Integer(), nullable=True))
            
        # Add thumbnail_url column
        if not column_exists('watch_history', 'thumbnail_url'):
            batch_op.add_column(sa.Column('thumbnail_url', sa.String(255), nullable=True))
            
        # Add upload_date column
        if not column_exists('watch_history', 'upload_date'):
            batch_op.add_column(sa.Column('upload_date', sa.DateTime(), nullable=True))

    # Backfill data
    conn = op.get_bind()
    conn.execute(text("""
        UPDATE watch_history 
        SET category = COALESCE(category, 'Unknown'),
            view_count = COALESCE(view_count, 0),
            subscriber_count = COALESCE(subscriber_count, 0),
            thumbnail_url = COALESCE(thumbnail_url, ''),
            upload_date = COALESCE(upload_date, watch_date)
        WHERE category IS NULL
        OR view_count IS NULL
        OR subscriber_count IS NULL
        OR thumbnail_url IS NULL
        OR upload_date IS NULL
    """))


def downgrade():
    with op.batch_alter_table('watch_history') as batch_op:
        if column_exists('watch_history', 'upload_date'):
            batch_op.drop_column('upload_date')
        if column_exists('watch_history', 'thumbnail_url'):
            batch_op.drop_column('thumbnail_url')
        if column_exists('watch_history', 'subscriber_count'):
            batch_op.drop_column('subscriber_count')
        if column_exists('watch_history', 'view_count'):
            batch_op.drop_column('view_count')
        if column_exists('watch_history', 'category'):
            batch_op.drop_column('category')