"""Add progress tracking columns to task_metadata

This migration adds columns that might be missing from databases created
before these columns were added to the model. Uses batch mode for SQLite
compatibility.

Revision ID: 0002
Revises: 0001
Create Date: 2025-01-01

Columns Added:
=============
- progress_current (INTEGER): Current progress value (default: 0)
- progress_total (INTEGER): Total progress value (default: 0)
- progress_message (VARCHAR): Human-readable progress status message
- metadata_json (JSON): Flexible JSON metadata storage

Migration Notes:
===============
- Uses SQLite batch mode (table recreation) for ALTER TABLE operations
- This is necessary because SQLite has limited ALTER TABLE support
- On large tables (>10K rows), this may take several seconds
- All existing data is preserved during the batch operation

Downgrade Behavior:
==================
The downgrade() function removes the four progress columns added by this
migration. Core data in other columns (task_id, status, etc.) is preserved.
Any data stored in progress_current, progress_total, progress_message, or
metadata_json columns will be LOST on downgrade.
"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect

# revision identifiers, used by Alembic.
revision = "0002"
down_revision = "0001"
branch_labels = None
depends_on = None


def column_exists(table_name: str, column_name: str) -> bool:
    """Check if a column exists in a table."""
    bind = op.get_bind()
    inspector = inspect(bind)

    if not inspector.has_table(table_name):
        return False

    columns = {col["name"] for col in inspector.get_columns(table_name)}
    return column_name in columns


def upgrade():
    """Add progress tracking columns to task_metadata if they don't exist."""
    # Only proceed if the table exists
    bind = op.get_bind()
    inspector = inspect(bind)

    if not inspector.has_table("task_metadata"):
        # Table doesn't exist yet, will be created by initial migration
        return

    # Check which columns need to be added
    columns_to_add = []

    if not column_exists("task_metadata", "progress_current"):
        columns_to_add.append(
            ("progress_current", sa.Integer(), {"server_default": "0"})
        )

    if not column_exists("task_metadata", "progress_total"):
        columns_to_add.append(
            ("progress_total", sa.Integer(), {"server_default": "0"})
        )

    if not column_exists("task_metadata", "progress_message"):
        columns_to_add.append(
            ("progress_message", sa.String(), {"nullable": True})
        )

    if not column_exists("task_metadata", "metadata_json"):
        columns_to_add.append(("metadata_json", sa.JSON(), {"nullable": True}))

    if not columns_to_add:
        # All columns already exist
        return

    # Use batch mode for SQLite compatibility
    with op.batch_alter_table("task_metadata") as batch_op:
        for col_name, col_type, kwargs in columns_to_add:
            batch_op.add_column(sa.Column(col_name, col_type, **kwargs))


def downgrade():
    """Remove progress tracking columns from task_metadata."""
    bind = op.get_bind()
    inspector = inspect(bind)

    if not inspector.has_table("task_metadata"):
        return

    columns_to_remove = [
        "progress_current",
        "progress_total",
        "progress_message",
        "metadata_json",
    ]

    existing_columns = {
        col["name"] for col in inspector.get_columns("task_metadata")
    }

    with op.batch_alter_table("task_metadata") as batch_op:
        for col_name in columns_to_remove:
            if col_name in existing_columns:
                batch_op.drop_column(col_name)
