"""add_task_hierarchy_fields

Adds task hierarchy fields to executions table for plan-task relationships.
Adds parent_task_id, task_sequence, and task_metadata columns.

Revision ID: add_task_hierarchy
Revises: add_planning_performance_indexes
Create Date: 2026-01-21 16:00:00.000000

"""
from typing import Sequence, Union

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

# revision identifiers, used by Alembic.
revision: str = 'add_task_hierarchy'
down_revision: Union[str, Sequence[str], None] = 'b8g9h1i2j3k4'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """
    Add task hierarchy fields to executions table.
    """

    # Add parent_task_id column
    op.add_column(
        'executions',
        sa.Column('parent_task_id', sa.String(255), nullable=True)
    )

    # Add task_sequence column
    op.add_column(
        'executions',
        sa.Column('task_sequence', sa.Integer(), nullable=True)
    )

    # Add task_metadata column
    op.add_column(
        'executions',
        sa.Column('task_metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default='{}')
    )

    # Create index on parent_task_id
    op.create_index(
        'idx_executions_parent_task_id',
        'executions',
        ['parent_task_id']
    )

    # Create composite index on plan_execution_id and task_sequence for ordering
    # Only create if plan_execution_id column exists (added by add_plan_executions migration)
    conn = op.get_bind()
    result = conn.execute(sa.text("""
        SELECT column_name FROM information_schema.columns
        WHERE table_name = 'executions' AND column_name = 'plan_execution_id'
    """))
    if result.fetchone():
        op.create_index(
            'idx_executions_plan_sequence',
            'executions',
            ['plan_execution_id', 'task_sequence']
        )
    else:
        print("⚠️ Skipping idx_executions_plan_sequence - plan_execution_id column not found")

    print("✅ Added task hierarchy fields to executions table")
    print("✅ Created indexes for parent_task_id and plan_sequence")


def downgrade() -> None:
    """
    Remove task hierarchy fields from executions table.
    """

    # Drop indexes (conditionally for plan_sequence index)
    conn = op.get_bind()
    result = conn.execute(sa.text("""
        SELECT indexname FROM pg_indexes
        WHERE tablename = 'executions' AND indexname = 'idx_executions_plan_sequence'
    """))
    if result.fetchone():
        op.drop_index('idx_executions_plan_sequence', table_name='executions')

    op.drop_index('idx_executions_parent_task_id', table_name='executions')

    # Drop columns
    op.drop_column('executions', 'task_metadata')
    op.drop_column('executions', 'task_sequence')
    op.drop_column('executions', 'parent_task_id')

    print("✅ Removed task hierarchy fields from executions table")
