"""add task management fields

Revision ID: add_task_mgmt_fields
Revises: add_indexes_opt
Create Date: 2026-01-14

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'add_task_mgmt_fields'
down_revision = 'add_indexes_opt'
branch_labels = None
depends_on = None


def upgrade():
    # Add retry configuration fields
    op.add_column('tasks', sa.Column('retry_count', sa.Integer(), nullable=False, server_default='0'))
    op.add_column('tasks', sa.Column('max_retries', sa.Integer(), nullable=False, server_default='3'))
    
    # Add timestamp fields for task lifecycle
    op.add_column('tasks', sa.Column('started_at', sa.Float(), nullable=True))
    op.add_column('tasks', sa.Column('completed_at', sa.Float(), nullable=True))
    
    # Add user context fields
    op.add_column('tasks', sa.Column('user_id', sa.String(), nullable=True))
    op.add_column('tasks', sa.Column('project_id', sa.String(), nullable=True))
    
    # Add indexes for new fields
    op.create_index('idx_tasks_user_id', 'tasks', ['user_id'])
    op.create_index('idx_tasks_project_id', 'tasks', ['project_id'])
    
    # Note: deleted_at column already exists from previous migration


def downgrade():
    # Remove indexes
    op.drop_index('idx_tasks_project_id', table_name='tasks')
    op.drop_index('idx_tasks_user_id', table_name='tasks')
    
    # Remove columns
    op.drop_column('tasks', 'project_id')
    op.drop_column('tasks', 'user_id')
    op.drop_column('tasks', 'completed_at')
    op.drop_column('tasks', 'started_at')
    op.drop_column('tasks', 'max_retries')
    op.drop_column('tasks', 'retry_count')