"""service_name_map_and_default_task_handling

Revision ID: 843e2d3f46d5
Revises: c6df714edd16
Create Date: 2026-02-09 17:44:08.622651

"""
from alembic import op
import sqlalchemy as sa
from utils.constants import DEFAULT_SERVICE_NAME, UNMAPPED_TASK_ID


# revision identifiers, used by Alembic.
revision = '843e2d3f46d5'
down_revision = 'c6df714edd16'
branch_labels = None
depends_on = None

BATCH_SIZE = 100

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###

    # Add is_autocreated column to tasks table with default False
    op.add_column('tasks', sa.Column('is_autocreated', sa.Boolean(), nullable=False, server_default=sa.false()))

    # Add is_system_task column to tasks table with default False
    op.add_column('tasks', sa.Column('is_system_task', sa.Boolean(), nullable=False, server_default=sa.false()))

    op.create_table('service_name_task_mappings',
    sa.Column('service_name', sa.String(), nullable=False),
    sa.Column('task_id', sa.String(), nullable=False),
    sa.Column('created_at', sa.TIMESTAMP(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
    sa.ForeignKeyConstraint(['task_id'], ['tasks.id'], name='fk_service_name_task_mappings_task_id', ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('service_name')
    )
    op.create_index('idx_service_name_mapping_task', 'service_name_task_mappings', ['service_name', 'task_id'], unique=False)
    op.create_index(op.f('ix_service_name_task_mappings_task_id'), 'service_name_task_mappings', ['task_id'], unique=False)

    # Create __unmapped__ task (marked as system task)
    op.execute(f"""
        INSERT INTO tasks (id, name, created_at, updated_at, is_agentic, archived, is_autocreated, is_system_task, task_metadata)
        VALUES ('{UNMAPPED_TASK_ID}', '{DEFAULT_SERVICE_NAME}', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, true, false, false, true, NULL)
    """)

    # Create mapping for __unmapped__ service name → __unmapped__ task
    op.execute(f"""
        INSERT INTO service_name_task_mappings (service_name, task_id, created_at)
        VALUES ('{DEFAULT_SERVICE_NAME}', '{UNMAPPED_TASK_ID}', CURRENT_TIMESTAMP)
    """)

    # Backfill NULL task_ids in spans table (batched, 100 rows at a time)
    connection = op.get_bind()

    total_spans_updated = 0

    while True:
        result = connection.execute(sa.text(f"""
            UPDATE spans
            SET task_id = '{UNMAPPED_TASK_ID}'
            WHERE id IN (
                SELECT id FROM spans
                WHERE task_id IS NULL
                LIMIT {BATCH_SIZE}
            )
        """))

        rows_updated = result.rowcount
        total_spans_updated += rows_updated

        if rows_updated < BATCH_SIZE:
            break

    print(f"Backfilled {total_spans_updated} spans with __unmapped__ task")

    # Backfill NULL task_ids in trace_metadata table (batched, 100 rows at a time)
    total_traces_updated = 0

    while True:
        result = connection.execute(sa.text(f"""
            UPDATE trace_metadata
            SET task_id = '{UNMAPPED_TASK_ID}'
            WHERE trace_id IN (
                SELECT trace_id FROM trace_metadata
                WHERE task_id IS NULL
                LIMIT {BATCH_SIZE}
            )
        """))

        rows_updated = result.rowcount
        total_traces_updated += rows_updated

        if rows_updated < BATCH_SIZE:
            break

    print(f"Backfilled {total_traces_updated} traces with __unmapped__ task")

    # After backfilling, make task_id NOT NULL in both tables
    op.alter_column('spans', 'task_id',
                   existing_type=sa.String(),
                   nullable=False)

    op.alter_column('trace_metadata', 'task_id',
                   existing_type=sa.String(),
                   nullable=False)


def downgrade() -> None:
    # Get the __unmapped__ task_id before deleting
    connection = op.get_bind()

    # First, revert NOT NULL constraints
    op.alter_column('spans', 'task_id',
                   existing_type=sa.String(),
                   nullable=True)

    op.alter_column('trace_metadata', 'task_id',
                   existing_type=sa.String(),
                   nullable=True)

    # Revert task_ids back to NULL in spans table (batched)
    total_spans_reverted = 0

    while True:
        result = connection.execute(sa.text(f"""
            UPDATE spans
            SET task_id = NULL
            WHERE id IN (
                SELECT id FROM spans
                WHERE task_id = '{UNMAPPED_TASK_ID}'
                LIMIT {BATCH_SIZE}
            )
        """))

        rows_updated = result.rowcount
        total_spans_reverted += rows_updated

        if rows_updated < BATCH_SIZE:
            break

    print(f"Reverted {total_spans_reverted} spans to NULL task_id")

    # Revert task_ids back to NULL in trace_metadata table (batched)
    total_traces_reverted = 0

    while True:
        result = connection.execute(sa.text(f"""
            UPDATE trace_metadata
            SET task_id = NULL
            WHERE trace_id IN (
                SELECT trace_id FROM trace_metadata
                WHERE task_id = '{UNMAPPED_TASK_ID}'
                LIMIT {BATCH_SIZE}
            )
        """))

        rows_updated = result.rowcount
        total_traces_reverted += rows_updated

        if rows_updated < BATCH_SIZE:
            break

    print(f"Reverted {total_traces_reverted} traces to NULL task_id")

    # Step 4: Delete __unmapped__ task (CASCADE will remove mapping)
    connection.execute(sa.text(f"""
        DELETE FROM tasks WHERE id = '{UNMAPPED_TASK_ID}'
    """))

    # Step 5: Drop indexes and table
    op.drop_index(op.f('ix_service_name_task_mappings_task_id'), table_name='service_name_task_mappings')
    op.drop_index('idx_service_name_mapping_task', table_name='service_name_task_mappings')
    op.drop_table('service_name_task_mappings')

    # Step 6: Remove is_system_task column from tasks table
    op.drop_column('tasks', 'is_system_task')

    # Step 7: Remove is_autocreated column from tasks table
    op.drop_column('tasks', 'is_autocreated')
