"""Download jobs: parent/shard fields for distributed bots.

Revision ID: 0007
Revises: 0006
Create Date: 2026-08-09

"""
from __future__ import annotations

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

revision: str = "0007"
down_revision: Union[str, None] = "0006"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    if not inspector.has_table("download_jobs"):
        return
    cols = {c["name"] for c in inspector.get_columns("download_jobs")}

    with op.batch_alter_table("download_jobs") as batch:
        if "parent_download_job_id" not in cols:
            batch.add_column(
                sa.Column(
                    "parent_download_job_id",
                    sa.String(36),
                    sa.ForeignKey("download_jobs.id"),
                    nullable=True,
                )
            )
        if "shard_index" not in cols:
            batch.add_column(sa.Column("shard_index", sa.Integer(), nullable=True))
        if "shard_count" not in cols:
            batch.add_column(sa.Column("shard_count", sa.Integer(), nullable=True))

    # Padre orquestador no tiene worker asignado
    op.alter_column(
        "download_jobs",
        "worker_id",
        existing_type=sa.String(36),
        nullable=True,
    )

    existing_indexes = {ix["name"] for ix in inspector.get_indexes("download_jobs")}
    if "ix_download_jobs_parent" not in existing_indexes:
        op.create_index(
            "ix_download_jobs_parent",
            "download_jobs",
            ["parent_download_job_id"],
        )


def downgrade() -> None:
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    if not inspector.has_table("download_jobs"):
        return
    existing_indexes = {ix["name"] for ix in inspector.get_indexes("download_jobs")}
    if "ix_download_jobs_parent" in existing_indexes:
        op.drop_index("ix_download_jobs_parent", table_name="download_jobs")

    cols = {c["name"] for c in inspector.get_columns("download_jobs")}
    with op.batch_alter_table("download_jobs") as batch:
        if "shard_count" in cols:
            batch.drop_column("shard_count")
        if "shard_index" in cols:
            batch.drop_column("shard_index")
        if "parent_download_job_id" in cols:
            batch.drop_column("parent_download_job_id")

    op.alter_column(
        "download_jobs",
        "worker_id",
        existing_type=sa.String(36),
        nullable=False,
    )